PyTorchで重回帰分析

class Net(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Net, self).__init__()

        self.fc1 = nn.Linear(input_dim, 100)
        self.fc2 = nn.Linear(100, 10)
        self.fc3 = nn.Linear(10, output_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))

        return self.fc3(x)

model = Net(26, 1).to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# X_trainに学習データの説明変数、X_test学習データの目的変数にデータが入っています
epochs = 20
data_size = len(X_train)

for epoch in range(epochs):
    loss_sum = 0

    for num in range(0, data_size):
        x = X_train[num].to(device)
        y = Y_train[num].to(device)

        optimizer.zero_grad()
        output = model(x)
        loss = criterion(output.reshape(y.shape), y)
        loss.backward()
        optimizer.step()
        loss_sum += loss.item()

    print(f"Epoch: {epoch+1}/{epochs}, Loss: {loss_sum / len(X_train)}")