Minimal PyTorch test script

minimal.py

import torch
from torch import nn
import tqdm


class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.tensor(1.))
        self.b = nn.Parameter(torch.tensor(0.))

    def forward(self, x):
        return self.w * x + self.b


x = torch.linspace(0, 1, 100, dtype=torch.float32)
y = 5 * x + 1

inp = x + torch.normal(0, .1, size=x.shape)

model = Model()
loss_fn = nn.MSELoss()
opt = torch.optim.SGD(model.parameters(), lr=.01, momentum=.9)

t = tqdm.trange(1000)
for i in t:
    opt.zero_grad()
    y_hat = model(inp)
    loss = loss_fn(y_hat, y)
    t.set_postfix(loss=loss.item())
    loss.backward()
    opt.step()

print(model.w.item(), model.b.item())