device_of

class torch.cuda.device_of(obj)

上下文管理器，用于将当前设备更改为给定对象的设备。

您可以将张量和存储用作参数。如果给定对象未在 GPU 上分配，则此操作为空操作。

参数

obj (张量 或 存储) – 在选定设备上分配的对象。