If you set CUDA_VISIBLE_DEVICEs=1,3,4, your GPUs will me masked and internally available as cuda:0, cuda:1, and cuda:2.


That should be the default as seen in this line of code 51


.cuda() will push your model/tensors to the default device. to() is the new style and I would recommend using it, as it’s easier to write device-agnostic code with it in my opinion.