device_of

class torch.cuda.device_of(obj)

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

你可以同时使用张量和存储作为参数。如果给定的对象没有在GPU上分配，则这是一个无操作。

Parameters

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