PyTorch Training Loop — Missing zero_grad Causes Nonsense

The training loop is the core execution pattern in PyTorch — a cycle that repeats for every batch of data: clear gradients, compute predictions, compute loss, backpropagate, update weights. PyTorch keeps this explicit on purpose. You are never far from the mechanics of optimization.

That explicitness is the trade-off. You get full visibility into gradient flow, parameter updates, device placement, mixed precision, gradient clipping, and scheduling. The cost is that there is no place to hide sloppy thinking. The loop is short, but it is stateful, and the order of operations matters.

The failure pattern I see most often in real code reviews is not an exotic math bug. It is a copy-pasted tutorial loop with one small change in the wrong place. Someone forgets optimizer.zero_grad(). Someone validates without model.eval(). Someone logs the loss tensor itself instead of loss.item() and wonders why GPU memory keeps growing. The loop is simple. The discipline around it is what separates a model that trains cleanly from one that burns two days of GPU time to produce nonsense.

What Is Training Loop in PyTorch Explained and Why Does It Exist?

The PyTorch training loop exists because optimization is stateful, and PyTorch chooses to make that state visible rather than hide it behind a one-line fit call. Each batch passes through the same sequence: clear old gradients, run the forward pass, compute the loss, backpropagate, and step the optimizer. That looks repetitive because it is repetitive. Training is controlled repetition.

The reason this pattern matters is that gradients in PyTorch accumulate by default. Parameters remember their previous .grad values until you clear them. Autograd also records the operations from the forward pass so backward can traverse that graph in reverse. The loop is not just procedural boilerplate — it is how you manage that state correctly.

In 2026, the canonical loop usually includes a few production-grade upgrades even for ordinary models: optimizer.zero_grad(set_to_none=True) for slightly lower memory traffic, mixed precision with torch.autocast on supported GPUs, gradient clipping when the model is deep or unstable, and explicit validation blocks with model.eval() plus torch.no_grad(). If your model is compile-friendly, torch.compile can sit on top of the same loop structure without changing the fundamentals.

What does not change is the contract. The loop still answers the same four questions every iteration: what did the model predict, how wrong was it, how should the weights change, and did they actually change.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader


def train_one_epoch(model, loader, criterion, optimizer, device, scaler=None, max_grad_norm=1.0):
    model.train()
    running_loss = 0.0
    total_samples = 0

    use_amp = device.type == 'cuda'

    for inputs, labels in loader:
        inputs = inputs.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        # 1) Clear stale gradients from the previous iteration
        optimizer.zero_grad(set_to_none=True)

        # 2) Forward pass + loss calculation
        with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=use_amp):
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        # 3) Backward pass
        if scaler is not None and use_amp:
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)  # unscale before clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)

            # 4) Optimizer step
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)
            optimizer.step()

        batch_size = inputs.size(0)
        running_loss += loss.item() * batch_size
        total_samples += batch_size

    return running_loss / total_samples


@torch.no_grad()
def validate(model, loader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    for inputs, labels in loader:
        inputs = inputs.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        outputs = model(inputs)
        loss = criterion(outputs, labels)

        running_loss += loss.item() * inputs.size(0)
        preds = outputs.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)

    avg_loss = running_loss / total
    accuracy = correct / total
    return avg_loss, accuracy


# Example wiring
# model = MyClassifier().to(device)
# model = torch.compile(model)  # optional in PyTorch 2.x when the model is compile-friendly
# optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
# criterion = nn.CrossEntropyLoss()
# scaler = torch.amp.GradScaler('cuda', enabled=(device.type == 'cuda'))
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
#
# for epoch in range(10):
#     train_loss = train_one_epoch(model, train_loader, criterion, optimizer, device, scaler)
#     val_loss, val_acc = validate(model, val_loader, criterion, device)
#     scheduler.step()  # epoch-level scheduler: step after the epoch
#     lr = optimizer.param_groups[0]['lr']
#     print(f'Epoch {epoch + 1:02d} | train_loss={train_loss:.4f} | val_loss={val_loss:.4f} | val_acc={val_acc:.2%} | lr={lr:.2e}')

Experiment Tracking: Logging Metrics and Checkpoints Like an Engineer

A training loop that only prints to stdout is fine for a notebook. It is not enough for a team. Once a model matters, you need a trace of what happened: which code ran, which hyperparameters were used, what the learning rate was at each epoch, what checkpoint corresponded to the best validation metric, and when the run started to go sideways if it did.

The simple pattern is still the right one. Log epoch-level metrics to a durable store. Keep checkpoints in object storage or a mounted artifact directory. Store the checkpoint path alongside the metrics row rather than pretending those two systems are unrelated. They are part of the same training story.

The practical benefit is not just reporting. It is rollback and diagnosis. When somebody says the new model is worse, you should be able to answer with evidence: which run, which epoch, which checkpoint, what the validation curve looked like, and whether the learning rate schedule or data version changed. Without that, model debugging turns into folklore.

-- io.thecodeforge: epoch-level training metrics for experiment tracking
INSERT INTO io.thecodeforge.training_history (
    run_id,
    model_id,
    epoch_number,
    train_loss,
    val_loss,
    val_accuracy,
    learning_rate,
    checkpoint_path,
    created_at
) VALUES (
    'run_2026_04_19_001',
    'ForgeResNet50-v3',
    12,
    0.2841,
    0.3198,
    0.9142,
    0.000300,
    's3://forge-models/ForgeResNet50-v3/epoch_12.pt',
    CURRENT_TIMESTAMP
);

-- Example rollback query: fetch the best validation checkpoint for a run
SELECT epoch_number, checkpoint_path, val_loss, val_accuracy
FROM io.thecodeforge.training_history
WHERE run_id = 'run_2026_04_19_001'
ORDER BY val_accuracy DESC, val_loss ASC
LIMIT 1;

Containerizing the Forge Training Environment

Training jobs fail for boring reasons far more often than people admit: wrong CUDA runtime, mismatched drivers, DataLoader workers starved by tiny shared memory, and buffered logs that make a crashed container look idle for ten minutes. Docker does not solve those problems automatically, but it gives you one place to make them explicit.

For 2026-era PyTorch stacks, three things matter immediately. First, pin the framework and CUDA versions rather than using latest. Second, use unbuffered Python output so logs appear in real time in whatever runtime you use. Third, remember that DataLoader workers share memory through /dev/shm inside the container. If you spin up multiple workers without enough shared memory, you get hangs, worker exits, or mysterious throughput collapse.

The other trap is silent CPU fallback. Teams assume the container is using the GPU because the base image has CUDA in its name. That proves nothing. The host still needs the NVIDIA Container Toolkit, the container still needs to run with --gpus all, and your startup logs should still print torch.cuda.is_available() plus the device name. If you do not verify that, you can burn hours training on CPU and only discover it when the epoch time looks absurd.

# io.thecodeforge: Reproducible PyTorch training container
# Pin versions. Never use a floating 'latest' tag in training infrastructure.
FROM pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime

WORKDIR /app

ENV PYTHONUNBUFFERED=1
ENV PIP_NO_CACHE_DIR=1

# System deps commonly needed by vision and tabular training stacks
RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    curl \
    libgl1 \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN pip install -r requirements.txt

COPY . .

# Good startup hygiene: print environment info before training begins
CMD ["python", "-u", "train.py"]

# Example runtime command:
# docker run --gpus all --shm-size=2g -v $(pwd)/data:/app/data -v $(pwd)/checkpoints:/app/checkpoints forge-trainer:latest

Common Mistakes and How to Avoid Them

Most training-loop bugs come from state you forgot was stateful. Gradients persist until you clear them. Dropout and BatchNorm switch behavior based on mode. Loss tensors keep their graph unless you detach or convert them properly for logging. PyTorch is explicit about all of this, but explicit does not mean self-correcting.

The two mode switches people most often misuse are model.train() and model.eval(). These are not decorative. They change the behavior of real layers. Validation without model.eval() is not a small mistake; it changes the model you think you are measuring. The same goes for validation without torch.no_grad() — maybe the metrics are still numerically correct, but you pay the full memory cost of graphs you will never use.

The other class of mistake is device mismatch. PyTorch will never silently move your batch to match the model. If the model is on GPU and the inputs are on CPU, the first forward pass fails. That is good. What is less obvious is partial mismatch inside more complex training code — an auxiliary tensor created on CPU in the middle of loss calculation, or class weights left on CPU while logits are on GPU. The discipline is the same: decide the device once, move everything that participates in the computation there, and verify it early.

import torch


def validate_model(model, data_loader, criterion):
    device = next(model.parameters()).device
    model.eval()  # CRITICAL: switch Dropout / BatchNorm to inference behavior

    total_loss = 0.0
    total_correct = 0
    total_samples = 0

    with torch.no_grad():
        for inputs, labels in data_loader:
            inputs = inputs.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)

            outputs = model(inputs)
            loss = criterion(outputs, labels)

            total_loss += loss.item() * inputs.size(0)
            preds = outputs.argmax(dim=1)
            total_correct += (preds == labels).sum().item()
            total_samples += labels.size(0)

    avg_loss = total_loss / total_samples
    accuracy = total_correct / total_samples

    model.train()  # restore training mode for the next epoch
    return avg_loss, accuracy


# Common anti-patterns to avoid:
# 1) Forgetting model.eval() before validation
# 2) Forgetting torch.no_grad() in validation
# 3) Logging loss instead of loss.item()
# 4) Using outputs.data instead of outputs.argmax(dim=1)
# 5) Forgetting to switch back to model.train() before the next epoch