Simple linear regression in pytorch

pytorch_linear_regression.py

import torch
from torch import nn
import numpy as np

net = nn.Sequential(
    torch.nn.Linear(2, 1)
)

def fun(x1, x2):
    return 2 * x1 + 3 * x2 + 5

x = np.random.rand(20000, 2) + 10
x = (x - x.mean()) / x.std()
y = np.array(list(map(lambda pair: fun(pair[0], pair[1]), x)))

loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(net.parameters(), lr=0.001)

def train(model, x, y, num_epoch):
    losses = []
    for i in range(num_epoch):
        model.train(True)
        loss = loss_fn(model(x), y)
        losses.append(loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.train(False)
    return losses

features = torch.FloatTensor(x)
labels = torch.FloatTensor(y).reshape(-1, 1)
train(net, features, labels, 2000)
print(net(torch.tensor([1.0, 1.0])))