I have updated pytorch and tried to run again. After updating, I am getting the following error.

RuntimeError: arguments are located on different GPUs at /py/conda-bld/pytorch_1490979338030/work/torch/lib/THC/generic/THCTensorIndex.cu:390

Now, I guess the problem is: when I trained the model, I didn’t set any specific GPU and as a result, pytorch used all available GPUs and now when I am trying to load the trained model and run, its giving me this error.
Is my understanding correct? Isn’t it possible to add some support in Pytorch to avoid these kind of circumstances? OR, do we need to set (using CUDA_VISIBLE_DEVICES) the gpu devices while training and then load the trained model on the same GPU to do the testing?