Writing every parameter by hand stops scaling at once. torch.nn holds the layers, torch.optim the update rules:

import torch
from torch import nn

x = torch.randn(64, 3)
y = x.sum(dim=1, keepdim=True)
# y = x @ W.T + b, W and b inside
model = nn.Linear(3, 1)
optimizer = torch.optim.SGD(
    model.parameters(), lr=0.1)
loss_fn = nn.MSELoss()

for step in range(300):
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
print(model.weight.data)  # near [1, 1, 1]

model.parameters() hands the optimizer every tensor it has to update, however many layers the model has. optimizer.zero_grad() and optimizer.step() replace the manual clearing and stepping from the last lesson.

A model's learned numbers live in model.state_dict(), which torch.save writes to a file and load_state_dict reads back.

Write train_linear(x, y, steps=300): train an nn.Linear from the columns of x (shape (n, k)) to y (shape (n, 1)) with torch.optim.SGD at lr=0.1 and nn.MSELoss, and return the trained model. Write final_loss(model, x, y) returning the mean squared error as a float, without recording gradients.

The explanation opens here once every check passes.