PyTorch Training Loop

A PyTorch training loop connects the model, data, loss, and optimizer. Writing these steps explicitly makes it possible to inspect a batch and control when gradients are cleared, parameters are updated, and validation is run. This example uses multiclass classification on a small synthetic dataset. Run the body blocks in order.

nn.Module: parameters that register themselves

Subclassing nn.Module gives you parameter tracking, recursive device movement, training/evaluation modes, and state-dict save and load. After super().__init__(), assigning a submodule directly to an attribute registers it. nn.Sequential registers its layers and calls them in order. For a collection you iterate over yourself, use nn.ModuleList or nn.ModuleDict; an ordinary Python list or dictionary does not register the layers it contains. Call model(x) to use the module’s normal call path and hooks, rather than invoking model.forward(x) directly.

import torch
import torch.nn as nn

torch.manual_seed(7)

class MLP(nn.Module):
    def __init__(self, in_dim, hidden, num_classes, p_drop=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.LayerNorm(hidden),
            nn.GELU(),
            nn.Dropout(p_drop),
            nn.Linear(hidden, num_classes),
        )

    def forward(self, x):
        return self.net(x)            # returns logits, not probabilities

model = MLP(784, 256, 10)
print(sum(p.numel() for p in model.parameters()))
print([n for n, _ in model.named_parameters()][:3])
# 204042
# ['net.0.weight', 'net.0.bias', 'net.1.weight']

This example uses CrossEntropyLoss, which expects unnormalized logits. A batch of shape (N, 784) produces (N, 10) scores and a scalar mean loss. Applying softmax before this loss changes the objective because it treats those probabilities as logits.

Assign custom trainable tensors as nn.Parameter attributes so the module exposes them to an optimizer. Use register_buffer for non-parameter state that should move with .to(device) and, by default, appear in the state dict. Such state is normally created with requires_grad=False; registration alone does not disable autograd. An ordinary tensor attribute is not automatically moved or saved, so it can end up on a different device from the module.

How autograd computes the gradients used by this loop is explained in PyTorch Tensors and Autograd for People Who Wrote Backprop by Hand.

Dataset and DataLoader

A map-style Dataset reports its length and returns an example for an index. Here each item is a feature vector and one integer class label. The DataLoader selects indices and collates the returned items: its default collation stacks feature vectors into (N,784) and scalar labels into (N,). Variable-length examples may need a custom collate_fn to pad them or return a different batch structure.

from torch.utils.data import Dataset, DataLoader

class TabularDataset(Dataset):
    def __init__(self, X, y, transform=None):
        self.X, self.y, self.transform = X, y, transform

    def __len__(self):
        return len(self.y)

    def __getitem__(self, idx):
        x = self.X[idx]
        if self.transform is not None:
            x = self.transform(x.clone())
        return x, self.y[idx]

# CPU keeps this first run independent of accelerator availability.
device = torch.device("cpu")
generator = torch.Generator().manual_seed(11)
centers = torch.randn(10, 784, generator=generator)
def make_data(n):
    y = torch.randint(0, 10, (n,), generator=generator)
    X = centers[y] + 0.5 * torch.randn(n, 784, generator=generator)
    return X.float(), y.long()

X_train, y_train = make_data(256)
X_val, y_val = make_data(80)
train_loader = DataLoader(
    TabularDataset(X_train, y_train),
    batch_size=64,
    shuffle=True,          # randomize these independent training rows
    num_workers=0,
    pin_memory=False,
    drop_last=False,
)
val_loader = DataLoader(TabularDataset(X_val, y_val), batch_size=64,
                        shuffle=False, num_workers=0)
print(X_train.shape, y_train.shape, X_train.dtype, y_train.dtype)
print(len(train_loader), len(val_loader))
# torch.Size([256, 784]) torch.Size([256]) torch.float32 torch.int64
# 4 2

The data are noisy observations around ten shared class centers; training and validation rows are drawn separately. The last validation batch contains 16 examples and is retained. LayerNorm in this MLP does not require dropping a small final batch. Start with num_workers=0 for a simple script or notebook. For multiprocessing on Windows or macOS, put executable setup under if __name__ == "__main__": and keep dataset definitions importable. CUDA runs can use pinned host memory with non-blocking transfers; that is an optional performance setting, not needed for this CPU example.

A stochastic transform in __getitem__ can produce fresh augmentation on repeated visits. The transform receives a clone because indexing a stored row can return a view: an in-place transform should not permanently alter the stored dataset. Batch transforms, GPU augmentation, and offline augmentation are also valid designs. This example uses no augmentation, and item loading runs in the main process because num_workers=0. Training shuffling is suitable for these independent rows; a task requiring temporal or grouped order needs an appropriate sampler instead.

The loop, in the order that matters

An epoch is one pass through the training loader. train() and eval() change behavior such as dropout; no_grad() separately disables recording for autograd. These functions use multiclass index targets and an unweighted, mean-reduced cross-entropy with no ignored labels. Multiplying each batch mean by its size gives the loss sum under those assumptions, including a short final batch. Class weights or ignored labels require a different denominator. Binary classification and regression also need different prediction metrics.

The driver uses OneCycleLR, which advances after each optimizer update. Here 3 epochs with 4 batches each give 12 updates. The supplied lr is its maximum learning rate, not the initial rate: the default div_factor=25 starts at lr/25. The default policy also cycles AdamW’s first-moment coefficient. Other schedulers may advance per epoch or use validation metrics. Gradient clipping at 1.0 is an optional choice; its threshold limits the gradient norm, not AdamW’s final update norm.

def train_one_epoch(model, loader, criterion, optimizer, scheduler, device):
    model.train()                                  # enable dropout in this MLP
    total_loss, total_n = 0.0, 0

    for xb, yb in loader:
        xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)      # clear accumulated gradients
        logits = model(xb)                         # forward
        loss = criterion(logits, yb)               # loss on logits
        loss.backward()                           # backward
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True)
        last_used_lr = optimizer.param_groups[0]["lr"]
        optimizer.step()                           # update
        scheduler.step()                           # prepare the next scheduled values

        total_loss += loss.detach() * xb.size(0)   # stays on device
        total_n += xb.size(0)

    if total_n == 0:
        raise ValueError("training loader is empty")
    return (total_loss / total_n).item(), last_used_lr


@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()                                   # disable dropout in this MLP
    loss_sum, correct, n = 0.0, 0, 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        logits = model(xb)
        loss_sum += criterion(logits, yb) * xb.size(0)
        correct += (logits.argmax(dim=1) == yb).sum()
        n += xb.size(0)
    if n == 0:
        raise ValueError("evaluation loader is empty")
    return (loss_sum / n).item(), (correct / n).item()

The code clears old gradients, computes logits and loss, differentiates that loss, and updates parameters. Clearing after the forward but before backward can also be valid. The important dependency is that gradients for the current update remain available until the optimizer uses them. The following mistakes change that behavior or waste resources.

• Missing model.train() / model.eval(): dropout stays active during evaluation, or BatchNorm keeps updating its running statistics on validation data. The metrics then describe a different model mode than intended. BatchNorm is not used in this MLP, but the distinction matters when it is present.
• Missing zero_grad: old batch gradients remain and alter the intended update; their contributions can reinforce or cancel each other. Intentional gradient accumulation needs its own update and scaling plan.
• zero_grad between backward and optimizer.step: the current gradients are cleared before the optimizer uses them. Clearing after the optimizer step is a valid alternative.
• scheduler.step() before optimizer.step(): the first update uses an unintended learning rate, and PyTorch warns about it.
• Accumulating loss.item() in the loop: on a GPU, reading each scalar back to Python introduces synchronization. Accumulate the detached tensor and convert once.
• Missing @torch.no_grad() on evaluation: evaluation records unnecessary computation graphs. In this particular loop, adding undetached losses also retains graph references across batches, increasing memory use and potentially causing an out-of-memory error.

Suppose the two validation batches contain 64 and 16 examples with mean losses 0.2 and 0.8. The epoch mean is \((64\times0.2+16\times0.8)/80=0.32\). Averaging the two means directly gives 0.5 and overweights the short batch. The loop’s loss_sum and example count implement the weighted calculation.

The reported training loss combines predictions made while weights change and dropout is active. Validation uses one fixed set of weights in evaluation mode. Their difference alone is not an overfitting diagnosis. To compare the datasets under the same conditions, evaluate both loaders with the same frozen checkpoint, deterministic preprocessing, and evaluation settings. The extra pass has a cost, so use it when the comparison matters.

Loss functions and what they expect

Task	Criterion	Model output	Target
Multiclass	CrossEntropyLoss	logits \((N,C)\)	int64 indices \((N,)\)
Binary	BCEWithLogitsLoss	logits \((N,)\)	float \((N,)\)
Multi-label	BCEWithLogitsLoss	logits \((N,C)\)	float \((N,C)\)
Regression	MSELoss / HuberLoss	\((N,1)\)	\((N,1)\)

This example uses int64 class indices shaped \((N,)\). CrossEntropyLoss also accepts floating-point probability targets shaped \((N,C)\), including valid one-hot distributions. In regression, comparing \((N,1)\) predictions with \((N,)\) targets can broadcast into pairwise errors shaped \((N,N)\); PyTorch may warn, but execution can continue. Check pred.shape == target.shape where matching elementwise targets are intended.

Parameter groups

Optimizers accept parameter groups with different settings. Common policies include excluding biases and normalization parameters from weight decay, or using a lower learning rate for a pretrained backbone. The helper below implements the first policy for this MLP and is called by the training driver.

def make_optimizer(model, lr):
    decay, no_decay = [], []
    for param in model.parameters():
        if param.requires_grad:
            (no_decay if param.ndim < 2 else decay).append(param)
    return torch.optim.AdamW([
        {"params": decay, "weight_decay": 0.01},
        {"params": no_decay, "weight_decay": 0.0},
    ], lr=lr)

optimizer = make_optimizer(model, 3e-4)
for group in optimizer.param_groups:
    print(group["weight_decay"], sum(p.numel() for p in group["params"]))
# 0.01 203264
# 0.0 778

For this MLP, the one-dimensional parameters are linear biases and LayerNorm gain and shift. The rule is a model-specific heuristic: for example, LayerNorm((2, 3)) has two-dimensional parameters. For other architectures, inspect module types and the resulting groups instead of assuming rank identifies every normalization parameter.

Seeding, configuration, and checkpointing for runs like this are covered in Reproducible Training Pipelines: Seeds, Configs, and Checkpoints.

Run epochs and save the selected checkpoint

def fit(model, train_loader, val_loader, epochs, device, checkpoint_path, lr=3e-4):
    if epochs < 1 or len(train_loader) == 0 or len(val_loader) == 0:
        raise ValueError("positive epochs and nonempty loaders are required")
    model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = make_optimizer(model, lr)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=lr, total_steps=epochs * len(train_loader))

    best_acc, history = float("-inf"), []
    for epoch in range(epochs):
        train_loss, last_used_lr = train_one_epoch(model, train_loader, criterion,
                                     optimizer, scheduler, device)
        val_loss, val_acc = evaluate(model, val_loader, criterion, device)
        if not all(torch.isfinite(torch.tensor(value)).item()
                   for value in (train_loss, val_loss, val_acc)):
            raise ValueError("nonfinite epoch metric; inspect data, loss, and updates")
        history.append((epoch, train_loss, val_loss, val_acc, last_used_lr))

        if val_acc > best_acc:
            best_acc = val_acc
            torch.save({"model": model.state_dict(),
                        "epoch": epoch,
                        "val_acc": val_acc}, checkpoint_path)

        print(f"epoch {epoch}  train {train_loss:.4f}  "
              f"val {val_loss:.4f}  acc {val_acc:.4f}")
    return history

Each history row records the zero-based epoch, mean training loss, validation loss, validation accuracy, and the first parameter group’s learning rate used for the epoch’s final update. Both groups share the schedule here. Reading the scheduler after its step would instead record the value prepared for a subsequent update, including an unused value after training ends. The clipping check rejects a nonfinite gradient norm before the optimizer update.

The checkpoint stores weights and validation metadata for inference, not a complete training-resume state. Resuming also requires optimizer and scheduler state, configuration, and relevant random-generator and data-order state. The architecture must match the saved weights. The driver rejects nonfinite epoch metrics and saves the first valid result, even at zero accuracy. Later results replace it only when accuracy increases; a tie keeps the earlier checkpoint even if validation loss improves. Choose that selection rule before inspecting test results.

from pathlib import Path
from tempfile import TemporaryDirectory

with TemporaryDirectory() as directory:
    checkpoint_path = Path(directory) / "best.pt"
    history = fit(model, train_loader, val_loader, epochs=3, device=device,
                  checkpoint_path=checkpoint_path, lr=3e-4)
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
    model.load_state_dict(checkpoint["model"])
    val_loss, val_acc = evaluate(model, val_loader, nn.CrossEntropyLoss(), device)
    print("restored epoch", checkpoint["epoch"])
    print("restored validation", round(val_loss, 4), round(val_acc, 4))
    print("saved accuracy matches", abs(val_acc - checkpoint["val_acc"]) < 1e-7)
# epoch 0  train 2.1851  val 1.6358  acc 0.7125
# epoch 1  train 1.2142  val 0.7937  acc 1.0000
# epoch 2  train 0.6738  val 0.6691  acc 1.0000
# restored epoch 1
# restored validation 0.7937 1.0
# saved accuracy matches True

The driver leaves the model at its last epoch. This example explicitly restores the best saved weights before evaluation. It uses a temporary directory so repeated practice runs do not overwrite a checkpoint; use a persistent path to retain it. Epoch losses and accuracy depend on the environment. Successful fitting on these deliberately easy synthetic classes checks the pipeline, not performance on a real application. The restored validation score is still a selection score; final assessment needs an untouched test set.

A single-batch fitting check can narrow down a training problem. Disable dropout and augmentation, then inspect whether the loss decreases on a fixed batch. Failure warrants checking labels, loss, gradients, and optimizer connections, but can also reflect learning rate, training duration, regularization, or capacity. Identical inputs with conflicting targets cannot all be memorized by a deterministic classifier.

Mixed precision, gradient accumulation, and profiling for this are covered in GPU Training: Mixed Precision, Gradient Accumulation, and Profiling.

References

• OneCycleLR
• ModuleList
• PyTorch. Saving and Loading Models.