The Five-Step Anatomy

Every PyTorch training step is the same five lines

Memorize this rhythm. You'll write it ten thousand times across your career, and it never changes:

optimizer.zero_grad() — clear accumulated gradients from the previous step.
output = model(batch) — forward pass. Builds the autograd graph.
loss = criterion(output, target) — compute the scalar loss.
loss.backward() — backward pass. Populates .grad on every leaf.
optimizer.step() — apply the update rule using the gradients.

The order is non-negotiable. Forward must happen before loss; loss before backward; backward before step; step after gradients exist; zero_grad before any of it (or at the very end of the previous iteration). Mixing them produces silent bugs.

What wraps the five steps

An outer loop over epochs (one full pass through the dataset).
An inner loop over batches from a DataLoader.
model.train() at the top of each epoch (or once before the loop, since training is the default mode).
A periodic evaluation pass (use model.eval() + torch.inference_mode()).

That's the whole skeleton. Everything else — schedulers, gradient clipping, mixed precision, logging, checkpointing — is a layer on top of these five steps.

Code

The minimal complete training loop·python

import torch
import torch.nn as nn

model = nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
criterion = nn.MSELoss()

for epoch in range(num_epochs):
    model.train()
    for x_batch, y_batch in train_loader:
        # 1. clear gradients
        optimizer.zero_grad()

        # 2. forward
        out = model(x_batch)

        # 3. loss
        loss = criterion(out, y_batch)

        # 4. backward
        loss.backward()

        # 5. step
        optimizer.step()

Common additions: clipping, scheduling, eval·python

import torch
import torch.nn as nn

# Setup
model = nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
criterion = nn.MSELoss()

for epoch in range(50):
    model.train()
    for x, y in train_loader:
        optimizer.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # ← optional but common
        optimizer.step()

    scheduler.step()                    # ← per-epoch scheduler

    # Quick eval pass
    model.eval()
    with torch.inference_mode():
        val_loss = sum(criterion(model(x), y).item() for x, y in val_loader) / len(val_loader)
    print(f"epoch {epoch}: val_loss = {val_loss:.4f}")

The 'safety unit test' — does my loop actually train?·python

import torch
import torch.nn as nn

# Tiny sanity test: can we overfit a single batch?
model = nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)
criterion = nn.MSELoss()

x = torch.randn(8, 10)
y = torch.randn(8, 2)

losses = []
for step in range(200):
    optimizer.zero_grad()
    loss = criterion(model(x), y)
    loss.backward()
    optimizer.step()
    losses.append(loss.item())

print(f"loss start: {losses[0]:.6f}, end: {losses[-1]:.6f}")
# loss start: 1.234567, end: 0.000012
# If end loss isn't << start loss, something is wrong with your loop.
# This 'overfit one batch' test catches 80% of training-loop bugs.

Every PyTorch training step is the same five lines

What wraps the five steps

Code

External links

Exercise

Progress

Comments 0