device_of

class torch.cuda.device_of(obj)

將當前裝置更改為給定物件裝置的上下文管理器。

您可以將張量和儲存作為引數。如果給定的物件未在 GPU 上分配，則此操作無效。

引數

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