Hi. Thanks for the replies. I wasn’t having much luck with gpu 2, so i killed the process running on gpu 0 to see if that would free up the memory for the job. Got this error:
RuntimeError: arguments are located on different GPUs at /opt/conda/conda-bld/pytorch_1535493744281/work/aten/src/THC/generated/../generic/THCTensorMathPointwise.cu:272

So I finnaly just allowed the process to find the default gpu:
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
and that uses the default gpu 0 and it works. BTW, torch version is 0.4.1.post2.
Thanks.