device_of

class torch.cuda.device_of(obj)

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

您可以使用张量（tensor）和存储（storage）作为参数。如果给定的对象未分配在 GPU 上，则此操作无任何影响。

参数:

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