device_of¶
-
class
torch.cuda.
device_of
(obj)[source]¶ 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.