PyTorch дает неправильные результаты из-за трансляции

Я хочу провести несколько экспериментов с нейронными сетями с PyTorch, но минимальный тестовый пример дает неправильные ответы. Тестовый пример устанавливает простую нейронную сеть с двумя входными переменными и выходной переменной, которая представляет собой просто сумму входных данных, и пытается изучить ее как задачу регрессии; Я ожидаю, что он сойдется при нулевой среднеквадратической ошибке, но на самом деле он сходится к 0,165. Вероятно, это связано с проблемой, упомянутой в предупреждающем сообщении; как я могу это исправить?

Код:

import torch
import torch.nn as nn

# data
Xs = []
ys = []
n = 10
for i in range(n):
    i1 = i / n
    for j in range(n):
        j1 = j / n
        Xs.append([i1, j1])
        ys.append(i1 + j1)

# torch tensors
X_tensor = torch.tensor(Xs)
y_tensor = torch.tensor(ys)

# hyperparameters
in_features = len(Xs[0])
hidden_size = 100
out_features = 1
epochs = 500

# model
class Net(nn.Module):
    def __init__(self, hidden_size):
        super(Net, self).__init__()
        self.L0 = nn.Linear(in_features, hidden_size)
        self.N0 = nn.ReLU()
        self.L1 = nn.Linear(hidden_size, 1)

    def forward(self, x):
        x = self.L0(x)
        x = self.N0(x)
        x = self.L1(x)
        return x


model = Net(hidden_size)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

# train
print("training")
for epoch in range(1, epochs + 1):
    # forward
    output = model(X_tensor)
    cost = criterion(output, y_tensor)

    # backward
    optimizer.zero_grad()
    cost.backward()
    optimizer.step()

    # print progress
    if epoch % (epochs // 10) == 0:
        print(f"{epoch:6d} {cost.item():10f}")
print()

output = model(X_tensor)
cost = criterion(output, y_tensor)
print("mean squared error:", cost.item())

Выход:

training
C:\Users\russe\Anaconda3\envs\torch2\lib\site-packages\torch\nn\modules\loss.py:445: UserWarning: Using a target size (torch.Size([100])) that is different to the input size (torch.Size([100, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
  return F.mse_loss(input, target, reduction=self.reduction)
    50   0.167574
   100   0.165108
   150   0.165070
   200   0.165052
   250   0.165039
   300   0.165028
   350   0.165020
   400   0.165013
   450   0.165009
   500   0.165006

mean squared error: 0.1650056540966034

И сообщение:

Предупреждение пользователя: использование целевого размера (torch.Size([100])), который отличается от входного размера (torch.Size([100, 1])). Это, вероятно, приведет к неправильным результатам из-за трансляции. Пожалуйста, убедитесь, что они имеют одинаковый размер.