Setting up a from-scratch training loop

Every framework hides the same three lines behind an abstraction. Once in a while it’s worth writing them out by hand — a forward pass, a loss, and a step — just to remember what the machinery is actually doing.

Here is the whole thing, minus the data loading:

for x, y in loader:
    logits = model(x)
    loss = cross_entropy(logits, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

That’s it. Everything else — schedulers, mixed precision, gradient clipping — is a refinement on top of these five lines. Call zero_grad() last (or first, but be consistent) and you’ve avoided the most common beginner mistake.

What to watch

Log the loss, sure, but also log the gradient norm. A loss that plateaus while the gradient norm quietly explodes is telling you something the loss curve alone never will.