How to set GPU ID? - Githubissues

raywzy commented 5 years ago

Very great package! My question is how to set the gpu device?

jeanfeydy commented 5 years ago

Hi @raywzy , Thanks! I have very little experience with multi-GPU configurations, but GeomLoss should run smoothly with the standard PyTorch interface:

gpuids = [0,1] if torch.cuda.device_count() > 1 else [0]
for gpuid in gpuids:
    with torch.cuda.device(gpuid):
        # Create your data on the gpu with .cuda(), work with GeomLoss, etc.

There is a short tutorial on the subject in the KeOps documentation. I cannot try it right now, but would expect GeomLoss to handle this transparently. Does it work on your machine?

Best regards, Jean

raywzy commented 5 years ago

@jeanfeydy That works fine for me. Thanks a lot :D

jeanfeydy commented 5 years ago

You're welcome :-)

jeanfeydy / geomloss

How to set GPU ID? #5