torch.set_default_device

torch.set_default_device(device) [source]

Устанавливает значение по умолчанию torch.Tensor для выделения на device. Это не влияет на вызовы фабричных функций, которые вызываются с явным аргументом device. Вызовы фабрики будут выполняться так, как если бы им был передан аргумент device.

Чтобы временно изменить значение по умолчанию устройства, а не устанавливать его глобально, используйте with torch.device(device): вместо этого.

Устройство по умолчанию изначально cpu. Если вы установите устройство тензора по умолчанию на другое устройство (например, cuda) без индекса устройства, тензоры будут выделены на том устройстве, которое является текущим для типа устройства, даже после вызова torch.cuda.set_device().

Предупреждение

Эта функция накладывает небольшую нагрузку на производительность при каждом вызове Python API torch (а не только для фабричных функций). Если это вызывает проблемы, оставьте комментарий на https://github.com/pytorch/pytorch/issues/92701

Параметры

device (device или строка) – устройство, которое нужно установить по умолчанию

Пример:

>>> torch.tensor([1.2, 3]).device
device(type='cpu')
>>> torch.set_default_device('cuda')  # current device is 0
>>> torch.tensor([1.2, 3]).device
device(type='cuda', index=0)
>>> torch.set_default_device('cuda:1')
>>> torch.tensor([1.2, 3]).device
device(type='cuda', index=1)