What is the standard PyTorch training loop and what does each step do?

The PyTorch training loop follows a fixed five-step pattern repeated for every batch. Understanding exactly what each line does — and what happens if you skip or reorder a step — is essential for debugging training issues.

import torch
import torch.nn as nn
import torch.optim as optim

model     = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 2))
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
loss_fn   = nn.CrossEntropyLoss()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def train_one_epoch(model, loader, optimizer, loss_fn, device):
    model.train()                       # 0. enables Dropout, BatchNorm train mode
    total_loss = 0.0

    for X_batch, y_batch in loader:
        X_batch = X_batch.to(device)
        y_batch = y_batch.to(device)

        optimizer.zero_grad()           # 1. clear gradients from previous step
        logits = model(X_batch)         # 2. forward pass
        loss   = loss_fn(logits, y_batch)  # 3. compute loss
        loss.backward()                 # 4. backpropagate — fills .grad
        optimizer.step()                # 5. update weights using gradients

        total_loss += loss.item() * X_batch.size(0)

    return total_loss / len(loader.dataset)

@torch.no_grad()                       # disable gradient tracking for eval
def validate(model, loader, loss_fn, device):
    model.eval()                        # disables Dropout, BatchNorm uses running stats
    total_loss, correct = 0.0, 0

    for X_batch, y_batch in loader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)
        logits = model(X_batch)
        loss   = loss_fn(logits, y_batch)
        total_loss += loss.item() * X_batch.size(0)
        correct    += (logits.argmax(1) == y_batch).sum().item()

    return total_loss / len(loader.dataset), correct / len(loader.dataset)

# Full training loop
for epoch in range(10):
    train_loss        = train_one_epoch(model, train_loader, optimizer, loss_fn, device)
    val_loss, val_acc  = validate(model, val_loader, loss_fn, device)
    print(f"Epoch {epoch}: train_loss={train_loss:.4f} val_loss={val_loss:.4f} val_acc={val_acc:.4f}")