Training Loop

The Training Loop

Training is a loop of five steps, repeated for each example:

1. Pick a training example
2. Forward pass: run the model, compute the loss
3. Backward pass: loss.backward() to compute all gradients
4. Update: nudge each parameter in the direction that reduces loss
5. Zero gradients: reset p.grad = 0 before the next step

Why Zero Gradients?

Gradients accumulate (we use +=) so they must be explicitly reset after each update. Otherwise, old gradients from previous examples would corrupt the current update.

Gradient Descent Update

The simplest optimizer just subtracts a fraction of the gradient:

def sgd_step(params, lr):
    for p in params:
        p.data -= lr * p.grad
        p.grad = 0

Example: Learning a Single Scalar

The clearest possible training loop: find w such that w * 3 = 3 (i.e., w = 1).

w = Value(0.0)
for step in range(20):
    pred = w * Value(3.0)
    loss = (pred - Value(3.0)) ** 2
    loss.backward()
    sgd_step([w], lr=0.05)

Each step, the gradient d(loss)/dw = 2 * (3w - 3) * 3 = 18(w-1) pushes w toward 1.

Your Task

Implement sgd_step(params, lr) and run the training loop.