"""Repeat the update on three examples until the model fits them.

Run: python3 examples/pytorch_training_loop.py
The browser lesson's training animation uses the same numbers.
"""

import torch
from torch import nn


# Three examples, one per row. The answers follow 2 × input + 1.
inputs = torch.tensor([[1.0], [2.0], [3.0]])
answers = torch.tensor([[3.0], [5.0], [7.0]])

model = nn.Linear(1, 1)
with torch.no_grad():
    model.weight.fill_(1.0)
    model.bias.fill_(0.0)

loss_fn = nn.MSELoss()  # The average of the three squared misses.
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for step in range(500):
    guesses = model(inputs)
    loss = loss_fn(guesses, answers)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step in (0, 9, 49, 499):
        print(f"Step {step + 1:3d}: loss {loss.item():.4f}  "
              f"weight {model.weight.item():.3f}  bias {model.bias.item():.3f}")

with torch.no_grad():
    print("A new input, 10, gives", round(model(torch.tensor([[10.0]])).item(), 2))