The Training Loop

Every PyTorch program trains with the same five lines, repeated once per round, an epoch. The model guesses with model(X). loss_fn(pred, y) gives back one number that says how far the guesses are from the answers. optimizer.zero_grad() clears the slopes left over from the last round, loss.backward() works out a new slope for every parameter, and optimizer.step() moves every parameter a small step against its slope. In your from-scratch network you wrote the slopes and the update by hand: here PyTorch does both.

Written once inside a function train(model, X, y, epochs), the five lines train any model on any data. Only the call changes:

model2 = nn.Linear(2, 1)     # a model with two inputs
train(model2, X2, y2, 51)    # the same train, other data, 51 epochs

Your train runs 101 epochs and prints only when epoch % 20 == 0, so you see epochs 0, 20, 40, 60, 80 and 100. loss.item() gives back the loss as a plain Python number, ready for round. The data is y = 3x + 2 with a little noise, so the program should print:

epoch 0 loss: 19.229
epoch 20 loss: 0.214
epoch 40 loss: 0.18
epoch 60 loss: 0.18
epoch 80 loss: 0.18
epoch 100 loss: 0.18
w: 2.9 b: 2.0

The check then calls your train again on a model with two inputs, and it must find y = 2 * x1 - x2 + 1: w: [2.0, -1.0] b: 1.0.

Your task: the data, the model, the prints and, inside train, the loss and the optimizer are ready. Inside the loop of train, above the if, write the five lines: the predictions in pred, the loss in loss, then clear the slopes, work out new ones, and take the step.

Assignment
Follow the Coding Tutorial to write the five lines of the training loop inside train.

Hint
The order matters: the loss needs pred, loss.backward() needs the loss, and optimizer.step() needs the slopes from backward(). All five lines sit eight spaces in, like the if.