device_of

class torch.cuda.device_of(obj)

上下文管理器，將目前裝置變更為給定物件的裝置。

您可以使用張量 (tensor) 和儲存 (storage) 作為參數。如果給定的物件未分配在 GPU 上，則此操作無效。

參數

obj (Tensor 或 Storage) – 分配在選定裝置上的物件。