Is there a way to specify the GPU to use when using a DataLoader with pin_memory = True?
I’ve tried using
torch.cuda.set_device(device)
and
with torch.cuda.device(gpu_id):
#DataLoader
but nvidia-smi shows that the GPU used is always the first available.
I think I could use:
CUDA_VISIBLE_DEVICES=x
to make the target GPU be the only one available in the notebook but I’d rather not.