Looking from the imagenet example it looks like multi-gpu training is pretty simple. All you need to do is add
net = torch.nn.DataParallel(net,device_ids=[0,1,2,3])
and you are good to go. Is this correct or I need to do something else also ?
I did this and I am getting this error
RuntimeError: cuda runtime error (10) : invalid device ordinal at torch/csrc/cuda/Module.cpp:84
CUDA_VISIBLE_DEVICES is properly set. Can someone tell me the fix ?