device_of

class torch.cuda.device_of(obj) [source]

Контекстный менеджер, который изменяет текущее устройство на устройство заданного объекта.

Вы можете использовать как тензоры, так и хранилища в качестве аргументов. Если заданный объект не выделен на графическом процессоре, это действие ничего не делает.

Параметры

obj (Тензор или Хранилище) – объект, выделенный на выбранном устройстве.