Building a PyTorch Training Loop with nn.Module, Dataset, and DataLoader

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 automatic parameter tracking, recursive device movement, mode switching between train and eval, and state-dict save and load. Any submodule assigned to self in __init__ is discovered automatically.

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()))     # 204042
print([n for n, _ in model.named_parameters()][:3])

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 implements __len__ and __getitem__ and returns one example. The DataLoader handles batching, shuffling, parallel loading, and collation.

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)
        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,          # True for training, False for validation
    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)
# torch.Size([256, 784]) torch.Size([256]) torch.float32 torch.int64
print(len(train_loader), len(val_loader))
# 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. Batch transforms, GPU augmentation, and offline augmentation are also valid designs. This dataset uses no augmentation; with num_workers=0, item loading runs in the main process.

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, so the scheduler advances after each optimizer update. Other schedules can operate per epoch or depend on validation metrics. Gradient-norm clipping at 1.0 is an optional choice in this example, not a required part of every training loop.

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)      # 2. clear accumulated grads
        logits = model(xb)                         # 3. forward
        loss = criterion(logits, yb)               # 4. loss on logits
        loss.backward()                            # 5. backward
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()                           # 6. update
        scheduler.step()                           # 7. per-batch schedule

        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()


@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()                                   # disable dropout, freeze BN
    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()

Each numbered line has a failure mode attached to it.

• 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.

Loss functions and what they expect

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.

A driver that leaves you something to look at

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 = train_one_epoch(model, train_loader, criterion,
                                     optimizer, scheduler, device)
        val_loss, val_acc = evaluate(model, val_loader, criterion, device)
        history.append((epoch, train_loss, val_loss, val_acc,
                        scheduler.get_last_lr()[0]))

        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

The checkpoint stores weights and validation metadata for inference, not a complete training-resume state. Resuming the same schedule would also require optimizer and scheduler states, along with relevant configuration and random-generator state. Saving a state dict avoids pickling the whole model object, whose loading is sensitive to class definitions and import paths; the model architecture still needs to match the saved weights. Starting best_acc at negative infinity ensures the first finite validation result is saved even when accuracy is zero.

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)

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.

Exercises

1. Order the six lines. Given zero_grad, forward, loss, backward, optimizer.step, scheduler.step, write the correct order. Then state what breaks for each of these three wrong orderings: zero_grad between backward and optimizer.step; scheduler.step before optimizer.step; loss computed after backward.

You should get: one correct ordering and three distinct failure modes.

2. Parameter groups. Split a model’s parameters into decay and no-decay groups using the ndim < 2 rule, then report how many parameters land in each and what fraction of the total the no-decay group holds.

You should get: a no-decay group with many tensors but a tiny share of the parameter count.

import torch.nn as nn
m = nn.Sequential(nn.Linear(784, 512), nn.LayerNorm(512), nn.GELU(),
                  nn.Linear(512, 256), nn.LayerNorm(256), nn.GELU(),
                  nn.Linear(256, 10))
decay = [p for p in m.parameters() if p.ndim >= 2]
no_decay = [p for p in m.parameters() if p.ndim < 2]
nd = sum(p.numel() for p in no_decay); d = sum(p.numel() for p in decay)
print(len(decay), len(no_decay), d, nd, f"{nd/(nd+d):.2%}")

3. Overfit one batch. Build the smallest useful test: take one batch of 32, turn off dropout and augmentation, and train until the loss is near zero. State what you would investigate if it fails, and why running it before any long training saves time.

You should get: a decreasing loss in this example; a failure is a reason to investigate, not proof of a particular cause.

import torch, torch.nn as nn
torch.manual_seed(0)
m = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 5))
x, y = torch.randn(32, 20), torch.randint(0, 5, (32,))
opt = torch.optim.Adam(m.parameters(), lr=1e-2)
crit = nn.CrossEntropyLoss()
for step in range(300):
    opt.zero_grad(set_to_none=True)
    loss = crit(m(x), y); loss.backward(); opt.step()
    if step % 100 == 0: print(step, round(loss.item(), 6))
with torch.no_grad():
    final_loss = crit(m(x), y)
print("final", round(final_loss.item(), 8))