Функция cuda.set_device

Функция cuda.set_device устанавливает активное устройство GPU для текущего процесса. Первым параметром функция принимает индекс устройства (целое число). Индексация начинается с нуля, поэтому первое устройство имеет индекс 0. После вызова все последующие операции на GPU будут выполняться на указанном устройстве, если не указано иное.

Синтаксис

torch.cuda.set_device(device)

Пример

Давайте установим активным устройство с индексом 0 и проверим текущее устройство:

import torch

torch.cuda.set_device(0)
current_device = torch.cuda.current_device()
print(current_device)

Результат выполнения кода:

0

Пример

Давайте создадим тензор после установки устройства и проверим его расположение:

import torch

torch.cuda.set_device(0)
t = torch.tensor([1, 2, 3, 4, 5]).cuda()
print(t.device)

Результат выполнения кода:

device(type='cuda', index=0)

Пример

Давайте переключимся между несколькими устройствами и проверим изменение текущего устройства:

import torch

torch.cuda.set_device(0)
print(torch.cuda.current_device())

torch.cuda.set_device(1)
print(torch.cuda.current_device())

Результат выполнения кода:

0
1

Смотрите также

функцию is_available,
которая проверяет доступность CUDA
функцию device_count,
которая возвращает количество доступных GPU
функцию current_device,
которая возвращает индекс текущего устройства
функцию get_device_name,
которая возвращает название устройства по индексу