Устройство модели в PyTorch

Данные линейного слоя должны лежать на том же устройстве, что и вход. Метод to переносит модуль целиком. Устройство задают объектом device со строкой процессора:

import torch
import torch.nn as nn

cpu = torch.device('cpu')
print(cpu) # выведет cpu

Создадим блок, перенесём его на это устройство и прочитаем место хранения веса:

import torch
import torch.nn as nn

class TinyNet(nn.Module):
	def __init__(self):
		super().__init__()
		self.layer = nn.Linear(2, 3)

	def forward(self, x):
		return self.layer(x)

net = TinyNet()
cpu = torch.device('cpu')
net = net.to(cpu)
weight = next(net.parameters())
print(weight.device) # выведет cpu

Создайте модуль с линейным слоем 2 на 3. Перенесите его на устройство процессора и выведите, где лежит первый тензор из перебора параметров.

Опишите класс с линейным слоем на три входа и один выход. Задайте устройство со строкой процессора, перенесите на него объект модели и выведите устройство любого параметра блока.

Соберите модуль с линейным слоем 4 на 4. После переноса на процессор выведите метку устройства у матрицы весов слоя.