Hello,
I am using PyTorch on a GPU cluster. When we send a job, the job runs on a machine with 4 GPUs. In the code, I do torch.device('cuda'). The problem is that this uses the GPU with id 0, but sometimes it’s the second or third GPU which is available.
I’d like my code to choose the first available GPU. I could iterate over torch.device('cuda:1'), torch.device('cuda:2'), etc. but it’s not very clean.
Do you have solutions ?