torch.set_default_device

torch.set_default_device(device) [исходный код]

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

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

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

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

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

Примечание

Это не влияет на функции, создающие тензоры, которые используют ту же память, что и входные данные, например: torch.from_numpy() и torch.frombuffer()

Параметры:

device (torch.device, str, int или None) – устройство, которое нужно задать по умолчанию, или None для сброса переопределения. Целое число интерпретируется как индекс текущего ускорителя.

Пример:

>>> torch.get_default_device()
device(type='cpu')
>>> torch.set_default_device('cuda')  # current device is 0
>>> torch.get_default_device()
device(type='cuda', index=0)
>>> torch.set_default_device('cuda')
>>> torch.cuda.set_device('cuda:1')  # current device is 1
>>> torch.get_default_device()
device(type='cuda', index=1)
>>> torch.set_default_device('cuda:1')
>>> torch.get_default_device()
device(type='cuda', index=1)