A Loss and an Optimizer

Last lesson you wrote the mean squared error yourself, moved w and b inside torch.no_grad(), and set their gradients back to zero. PyTorch has ready-made pieces for these jobs, so every training loop looks the same.

nn.MSELoss() gives back a loss function. Call it with the predictions and the true answers, in that order, and it gives back one number: the mean of the squared errors.

loss_fn = nn.MSELoss()
loss = loss_fn(torch.tensor([3.0]), torch.tensor([1.0]))
print(loss.item())    # (3 - 1) squared: 4.0

torch.optim.SGD(model.parameters(), lr=0.1) gives back an optimizer that holds the model's weight, its bias and the learning rate. Each round you call three methods: optimizer.zero_grad() sets every .grad back to zero, loss.backward() fills them with the slopes, and optimizer.step() moves each parameter by lr times its gradient, the update you wrote by hand.

The data follows y = 3x + 2. After 60 rounds your code prints the line it found, and the check shows the loss is 0.0 and x = 10 gives 32.0:

w: 3.0 b: 2.0

Your task: x, y, the model nn.Linear(1, 1) and a loop that predicts are written for you. Create loss_fn and an SGD optimizer with lr=0.1, then inside the loop compute loss and call optimizer.zero_grad(), loss.backward() and optimizer.step().

Assignment
Follow the Coding Tutorial to train the line with nn.MSELoss and an SGD optimizer instead of the hand-written loss and update.

Hint
Order matters inside the loop: clear the gradients, then backward(), then step(). step() uses the slopes that backward() just made.