import torch

x = torch.tensor([1., 2., 3., 4., 5.])
y = torch.tensor([20., 40., 60., 80., 100.])
w = torch.tensor(0.0, requires_grad=True)
b = torch.tensor(0.0, requires_grad=True)
lr = 0.001

for epoch in range(1000):
    prediction = x * w + b
    loss = torch.mean((prediction - y) ** 2)
    w.grad = None
    b.grad = None
    loss.backward()
    if epoch % 100 == 0:
        print(f"epoch={epoch:4d} loss={loss.item():.3f} w={w.item():.3f} b={b.item():.3f}")

    with torch.no_grad():
        w -= lr * w.grad
        b -= lr * b.grad

with torch.no_grad():
    print("Final MSE:", torch.mean((x * w + b - y) ** 2).item(), "w:", w.item(), "b:", b.item())
    print("Prediction for 7:", (7 * w + b).item())