What is the correct usage of torch.cuda.device(idx)?

1

What is the correct usage of torch.cuda.device ? Set_device has the comment “Usage of this function is discouraged in favor of device”. But when I try to use it to set the current device it doesn’t work, whereas set_device does work?

   (Pdb) torch.cuda.device_count()
    2
    (Pdb) torch.cuda.current_device()
    0
    (Pdb) torch.cuda.device(1)
    <torch.cuda.device object at 0x2afd1ce77390>
    (Pdb) torch.cuda.current_device()
    0
    (Pdb) torch.cuda.set_device(1)
    (Pdb) torch.cuda.current_device()
    1
2

Hi,

this should answer your question.

1 Like