device
class torch.cuda.device(device)
Context-manager that changes the selected device.
Parameters:
      device (torch.device or int) -- device index to
      select. It's a no-op if this argument is a negative integer or
      "None".