Running different processes on different GPUs

I am using python multiprocessing to spawn multiple processes which run on different model objects of their own. I also have multiple GPUs available with me. How can I allocate different GPUs to different processes(as in each model running on separate GPU)? Does Pytorch do this by default or does it run all processes on 1 GPU only unless specified?

PyTorch doesn’t do this by default. You need to specifically specify your desired GPU by .cuda(gpu_id).

Yeah I tried that. How can I see the available gpu id’s inside Pytorch only though?

Could this work

when i print torch.cuda.device_count() it shows 4. But when inside a process i do torch.cuda.set_device(1)(it works fine in a single process code), it gives multiprocesssing error(just tells the file name in traceback, nothing else)
Process Process-2:
Traceback (most recent call last):
File “/tools/anaconda3/envs/py35/lib/python3.5/multiprocessing/”, line 252, in
torch.cuda.set_device(0) works absolutely fine though.
Could you please provide any help in this regard?

Uhh, interesting. How are you doing mp?

@dakshanand you probably didn’t switch to spawn mode. See:

@SimonW, @smth

Line 21 - I am using torch.multiprocessing
Line 240 - I am using ‘spawn’ method
Line 162 - torch.cuda.set_device(1) [ If I put 0 here it works fine]

Could you put spawn line right after you import mp, and try again?