device_of
class torch.cuda.device_of(obj)
Context-manager that changes the current device to that of given
   object.
You can use both tensors and storages as arguments. If a given
   object is not allocated on a GPU, this is a no-op.
Parameters:
      obj (Tensor or Storage) -- object allocated on the
      selected device.