Train a Line by Hand in Python

Training is a loop. Every round, the model makes its predictions and measures the loss. Then loss.backward() finds the slopes, and each weight takes a small step against its slope: the slope times the learning rate lr, subtracted. In the from-scratch chapter you worked the slopes out yourself. Now backward() leaves them in w.grad and b.grad.

Two new tools keep the update clean. with torch.no_grad(): opens a block where PyTorch stops recording, so changing w and b isn't treated as part of the formula. Also, backward() adds to .grad, so every round ends by setting the slopes back to zero with .grad.zero_(). The _ at the end means the tensor changes in place. Here it is on another tensor:

v = torch.tensor(5.0, requires_grad=True)
(v ** 2).backward()      # v.grad is 10.0
with torch.no_grad():
    v -= 0.1 * v.grad    # v is now 4.0
v.grad.zero_()           # v.grad is 0.0 again

The five points lie on the line y = 3x + 2. Training starts from w = 0 and b = 0 with lr = 0.1. The loop that makes the predictions and the mean squared error is given, and it prints the loss every 10 rounds. On its own it prints 22.0 every time, because nothing changes w or b. Your finished program must print the lines below. The last three come from the lesson's check, which also trains a copy without the zeroing.

epoch 0 loss: 22.0
epoch 10 loss: 0.047
epoch 20 loss: 0.001
epoch 30 loss: 0.0
epoch 40 loss: 0.0
epoch 50 loss: 0.0
w: 3.0 b: 2.0
x = 10 -> 32.0
w.grad now: 0.0
without zero_(), epoch 50: loss 6.532 w: 3.484 b: 2.673

Your task: in the loop, below the if, call loss.backward(). Then move w and b against their gradients inside with torch.no_grad():, and set both gradients back to zero.

Assignment
Follow the Coding Tutorial to add the backward call, the update and the zeroing to the training loop.

Hint
Watch the indentation. backward() and the two zero_() lines sit four spaces in, inside the loop. The two updates sit eight spaces in, inside the with block. Use -=, so w and b stay the same tensors that PyTorch tracks.