The training loop — from the 5-line core to production

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

That's it. Every neural network training run — MNIST, ImageNet, GPT-5 — has this loop at its center. zero_grad clears the gradient buffer (autograd accumulates by default — without it your gradients sum across batches and the model trains on garbage). loss.backward() populates .grad on every parameter. optimizer.step() reads those grads and updates the parameters in place.

Everything beyond the five lines — mixed precision, gradient accumulation, distributed training, EMA weights, schedulers — is decorative. Understand the core, and you can read any training script ever written. Throughout this chapter we'll add one decoration at a time and explain why it earns its place.

• `nn.CrossEntropyLoss()` for multi-class — takes raw logits, applies log_softmax + NLL internally. Do not softmax before feeding it. Use weight=tensor([..]) for class imbalance, label_smoothing=0.1 for regularization.
• `nn.BCEWithLogitsLoss()` for binary / multi-label — takes raw logits, applies sigmoid + log internally. Numerically stable. Use pos_weight=tensor([..]) for imbalance.
• `nn.MSELoss()` for regression — squared error. Maximum-likelihood under Gaussian noise.
• `nn.L1Loss()` for robust regression — absolute error. Less sensitive to outliers.
• `nn.HuberLoss(delta=1)` — quadratic near zero, linear far away. Best of MSE + L1.
• `nn.CosineEmbeddingLoss` for similarity learning.
• `nn.TripletMarginLoss` for metric learning (anchor / positive / negative).
• `nn.KLDivLoss(reduction='batchmean')` for distillation — student log-probs vs. teacher probs.
• Custom contrastive losses (InfoNCE, NT-Xent) for self-supervised work.

The most common loss-related bug: feeding softmax probabilities into CrossEntropyLoss. It 'works' (no error) but trains poorly because you double-apply softmax. Memorize: logits → loss, never softmax before the loss layer.

SGD + momentum: default for vision CNNs trained from scratch. lr=0.1, momentum=0.9, weight_decay=5e-4. Combine with a step or cosine schedule.

AdamW: default for transformers, fine-tuning, and most modern work. lr=3e-4 from scratch, lr=2e-5 for fine-tuning pretrained. betas=(0.9, 0.999), weight_decay=0.01. Use eps=1e-8 (or 1e-6 for bf16 stability).

Lion: simpler than Adam, faster on some workloads (lr ~10× smaller than Adam).

Schedulers:

• LinearLR warmup → CosineAnnealingLR decay: the modern transformer recipe.
• OneCycleLR: fast.ai's super-convergence recipe.
• StepLR(gamma=0.1, step_size=30): classic vision recipe.
• ReduceLROnPlateau: monitors a metric, drops LR when it plateaus.

from torch.optim.lr_scheduler import SequentialLR, LinearLR, CosineAnnealingLR
warmup = LinearLR(opt, start_factor=0.01, end_factor=1.0, total_iters=warmup_steps)
decay = CosineAnnealingLR(opt, T_max=total_steps - warmup_steps, eta_min=1e-6)
sched = SequentialLR(opt, [warmup, decay], milestones=[warmup_steps])

Karpathy's rule: if you only have time to tune one thing, tune the learning rate. Warmup for the first few hundred steps stabilizes early gradients; cosine annealing to zero squeezes out the last bit of performance.

Mixed precision halves memory and ~2× throughput on modern GPUs.

scaler = torch.cuda.amp.GradScaler()    # only needed for fp16; bf16 skip it

for x, y in loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        logits = model(x)
        loss = loss_fn(logits, y)
    loss.backward()    # for bf16. fp16: scaler.scale(loss).backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()   # for bf16. fp16: scaler.step(opt); scaler.update()

Gradient clipping with clip_grad_norm_(params, max_norm=1.0) prevents exploding gradients. Essential for transformers, RNNs, and any model trained on noisy data. Apply *after* backward, *before* step.

Gradient accumulation lets you simulate a larger batch when memory is tight:

for step, (x, y) in enumerate(loader):
    with torch.autocast(...):
        loss = loss_fn(model(x), y) / accum_steps
    loss.backward()
    if (step + 1) % accum_steps == 0:
        clip_grad_norm_(...)
        optimizer.step()
        optimizer.zero_grad()
        scheduler.step()

Divide the loss by accum_steps to keep the effective gradient magnitude correct. Run optimizer step only every N micro-batches.

Checkpoints save you from crashes and let you resume:

state = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'scheduler': scheduler.state_dict(),
    'scaler': scaler.state_dict(),
    'epoch': epoch, 'step': step,
    'best_val': best_val,
}
torch.save(state, 'ckpt.pt')
# resume
state = torch.load('ckpt.pt')
model.load_state_dict(state['model'])
# ... and the rest

Save every N steps, keep the last K and the best-by-val-metric. Use torch.save(..., _use_new_zipfile_serialization=True) (the default in modern PyTorch).

Exponential Moving Average (EMA) of the weights gives a smoother, often better-generalizing model for inference:

from torch.optim.swa_utils import AveragedModel, get_ema_avg_fn
ema_model = AveragedModel(model, avg_fn=get_ema_avg_fn(0.999))
# in training loop, after optimizer.step():
ema_model.update_parameters(model)
# at eval: use ema_model.module instead of model

Diffusion models, GANs, and many SOTA vision models report EMA weights.

Early stopping: track validation loss, stop when it hasn't improved for patience=K epochs. Saves compute and reduces overfitting.

DistributedDataParallel (DDP) replicates the model on each GPU, syncs gradients via all-reduce after each backward, and updates locally. Linear scaling up to ~64 GPUs for most workloads. The pattern:

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
model = DDP(model.to(local_rank), device_ids=[local_rank])
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler, ...)

Launch with torchrun --nproc_per_node=8 train.py.

Fully Sharded Data Parallel (FSDP) — shards parameters, gradients, and optimizer state across GPUs (ZeRO-3). For models too big to fit on one GPU even at batch=1.

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model, auto_wrap_policy=transformer_auto_wrap_policy)

Mixed strategies: FSDP + tensor parallelism (Megatron) + pipeline parallelism (PP) is what trains GPT-scale models. For 99% of practitioners, DDP is sufficient; FSDP unlocks 7B-70B fine-tuning on commodity 8×A100/H100 nodes.