device_of

class torch.cuda.device_of(obj)

上下文管理器，將當前裝置切換到給定物件所在的裝置。

您可以使用張量和儲存作為引數。如果給定物件未在 GPU 上分配，則此操作無效（no-op）。

引數

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