Assemble full loop (data, forward, loss, backward, optimizer step)
The problem, from first principles
Every piece this question needs already exists somewhere earlier in this curriculum: [03-training-loop/01-dataset-dataloader]'s DataLoader produces batches, [02-layers/01-linear-forward]'s linear_forward and [02-layers/02-linear-backward]'s linear_backward handle the model's forward and backward pass, and [03-dl-training/01-optimizers]'s sgd_step (or any of the other optimizer questions) handles the parameter update. What's been MISSING is the glue that wires all of them together into the actual loop every training script runs: pull a batch, run it forward, measure how wrong the prediction was, backpropagate that error, nudge the parameters a little in the direction that reduces it, and repeat for every batch in the dataset.
This is, in a very real sense, the entire "training a model" story in miniature: everything more sophisticated later in this curriculum, deeper networks, attention, transformers, is still fundamentally this same five-step loop, just with a more elaborate forward pass and a more elaborate model to differentiate through.
From theory to code
Implement mse_loss_and_grad(pred, target), which returns (loss, grad_pred) together in one call (mean squared error, and its gradient with respect to pred), and train_one_epoch(loader, weight, bias, lr), which loops over every batch DataLoader yields, and for each one: runs linear_forward to get a prediction, calls mse_loss_and_grad to get the loss and its gradient, calls linear_backward to backpropagate that gradient into grad_weight and grad_bias, and updates weight/bias with a plain SGD step (param - lr * grad). Track the running total loss across batches and return the average, along with the updated weight and bias.
Constraints
mse_loss_and_gradcomputes MSE as the mean of squared errors over ALL entries (not just the batch dimension, ifpredhas multiple output features, average over those too).train_one_epochmust updateweightandbiasimmediately after each batch (not accumulate gradients across the whole epoch and update once at the end); this is standard mini-batch SGD, not full-batch gradient descent.- The returned average loss is the mean of each batch's OWN average loss across all batches in the epoch, not weighted differently by batch size.
- Reuse
linear_forwardandlinear_backward, don't reimplement the matrix operations here.
Hints
Hint 1: mse_loss_and_grad
loss = mean((pred - target)^2). For the gradient: since loss is the mean over pred.size total entries of (pred_i - target_i)^2, and d/dp[(p-t)^2] = 2(p-t), the gradient of the MEAN is 2*(pred - target) / pred.size (dividing by the total element count, not just the batch size, to match the mean in the loss).
Hint 2: One epoch, one batch at a time
for batch_x, batch_y in loader: gives you one batch per iteration. Inside the loop: pred = linear_forward(batch_x, weight, bias), then loss, grad_pred = mse_loss_and_grad(pred, batch_y), then _, grad_weight, grad_bias = linear_backward(grad_pred, batch_x, weight).
Hint 3: Updating and tracking
weight = weight - lr * grad_weight and bias = bias - lr * grad_bias, applied INSIDE the loop, right after computing the gradients for that batch (not saved up for later). Accumulate total_loss += loss and a batch counter, then return total_loss / n_batches after the loop ends.
Click "Run Tests" to test your implementation