Перенос на видеокарту в PyTorch

Метод to переносит тензор на указанное устройство. Если карты нет, в качестве цели берут процессор, а не отдельный метод только для видеокарты.

Сначала выберем устройство по результату проверки карты, затем перенесём маленький тензор и прочитаем метку устройства:

import torch

device = torch.device(
	"cuda" if torch.cuda.is_available() else "cpu"
)
row = torch.tensor([1.0, 2.0])
row = row.to(device)
print(row.device) # выведет cpu

Тот же приём подходит для модулей: их тоже переносят методом to, чтобы веса и вход лежали на одном устройстве.

Соберите таблицу [[1.0, 2.0], [3.0, 4.0]], перенесите её на устройство по проверке карты и выведите метку устройства.

Создайте объект устройства для процессора, перенесите на него число 5.0 и выведите метку устройства результата.

Соберите ряд из 10, 20 и 30, перенесите его туда же, куда ведёт проверка карты, и выведите метку устройства.