torch.cuda.set_device
torch.cuda.set_device(device)
Sets the current device.
Usage of this function is discouraged in favor of "device". In most
   cases it's better to use "CUDA_VISIBLE_DEVICES" environmental
   variable.
Parameters:
      device (torch.device or int) -- selected device. This
      function is a no-op if this argument is negative.