device_of

class torch.cuda.device_of(obj)

上下文管理器，将当前设备切换到给定对象所在的设备。

您可以使用张量和存储作为参数。如果给定对象未在 GPU 上分配，则此操作无效（no-op）。

参数

obj (Tensor 或 Storage) – 在选定设备上分配的对象。