If my model has no dropout or batch normalization layers, do I still need to call model.eval()?

Calling it is still good practice even though, for a model with genuinely no dropout or batch norm layers, it would have no observable effect right now. The value is in the habit: it costs nothing, and it protects you the moment you add a layer later (directly or through a library component) whose behavior does depend on the mode, without you having to remember to add the call at that point.

Can I load a checkpoint that was saved on a GPU machine onto a CPU-only machine?

Yes, but it needs an explicit `map_location='cpu'` argument passed to `torch.load` — without it, PyTorch tries to restore the tensors onto the same GPU device they were saved from, which fails outright on a machine with no GPU. This is the same underlying principle as device placement elsewhere in PyTorch: it is explicit, not inferred automatically.

Do I need multiple GPUs to use mixed precision training?

No — mixed precision is a single-GPU technique. It changes which numeric format operations run in on whichever GPU you already have; it has nothing to do with how many GPUs are involved. Multi-GPU data or model parallelism is a separate, complementary decision.

When should I reach for model parallelism instead of just using a GPU with more memory?

Model parallelism is worth the added complexity specifically when a single model's weights genuinely do not fit on the largest GPU you can reasonably get, which is common for large language models but rare otherwise. If a bigger single GPU (or mixed precision, or gradient accumulation) would solve the problem, those are simpler and should be tried first — model parallelism adds real communication overhead between GPUs that a single-GPU fix avoids entirely.

beginnerintermediateadvanced

AI Full-Stack · Section 10

PyTorch

Level: intermediate
Read: 55 min
Concepts: 6

Section 9's concepts made concrete in the framework most deep learning is actually written in: tensors, devices and autograd as the core vocabulary; the full training/validation loop assembly, with the mode-switching and no_grad bugs that most commonly break it; GPU memory as the usual practical limit and mixed precision as the standard fix; saving a model safely as a state_dict rather than a whole object; the two distributed-training strategies and which bottleneck each solves; and the debugging approach that fits PyTorch's three most common runtime errors — shape, device and dtype mismatches.

What is true here

A tensor and its model must share a device explicitly — PyTorch errors rather than silently copying
model.train()/model.eval() change real layer behavior (dropout, batch norm), not just bookkeeping
GPU memory, not compute speed, is usually the real limit — mixed precision roughly halves it
Save a state_dict, not the whole model object — it forces the architecture to be explicit at load time
Shape/device/dtype errors are hard runtime errors by design — read the exact values, then trace backward

What you will be able to do

Move tensors and models to the same device explicitly, and explain why PyTorch does not do it silently
Write a training loop that correctly switches between train() and eval() mode and uses no_grad() in validation
Explain why GPU memory, not speed, is usually the real constraint, and how mixed precision addresses it
Save and reload a model using its state_dict rather than the whole object
Debug a shape, device, or dtype error by reading the message and tracing backward to its source

Deep learning fundamentals in PyTorch

The framework vocabulary, the training loop assembly, and the practical debugging skills that turn section 9's mechanism into runnable, production-shaped code.

PyTorch tensors, devices, autograd, modules and dataloaders

standardintermediate

A PyTorch tensor is like a NumPy array, but it can live on a GPU and track the operations performed on it for automatic differentiation. A device specifies where a tensor's data actually lives (CPU or GPU). Autograd is the system that builds the computational graph (§9.1) and computes gradients automatically. A `Module` defines a model's layers; a `Dataset` and `DataLoader` handle loading and batching training data.

Think of it as

These five pieces are the vocabulary every PyTorch program is written in. A tensor is the core data structure — everything (inputs, weights, gradients) is a tensor. A device is an explicit choice: a tensor and the model operating on it must be on the same device, or PyTorch raises an error rather than silently moving data, because a silent cross-device copy would be a hidden performance cost. Autograd is what makes `loss.backward()` (§9.2) work — it only tracks gradients for tensors with `requires_grad=True`, which is why inference code wraps itself in `torch.no_grad()` to skip that (unneeded, memory-costly) bookkeeping. A `Module` is the base class every model architecture is built from — it groups layers and defines the forward pass; PyTorch tracks all its parameters automatically once they're assigned as attributes. A `Dataset` defines how to fetch one example; a `DataLoader` wraps it to handle batching, shuffling, and parallel loading, which is what actually produces the batches a training loop iterates over.

python

import torch
from torch.utils.data import Dataset, DataLoader

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

class MyDataset(Dataset):
    def __len__(self): return len(self.data)
    def __getitem__(self, i): return self.data[i], self.labels[i]

loader = DataLoader(MyDataset(), batch_size=32, shuffle=True)

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(10, 1)
    def forward(self, x):
        return self.fc(x)

model = MyModel().to(device)   # move model to the same device as its inputs

Broadcasting is real repetition, not magic

Remember: Tensors are the core data structure, device-aware and gradient-trackable via autograd. A tensor and its model must share a device explicitly. `nn.Module` defines architecture and tracks parameters; `Dataset` + `DataLoader` handle fetching and batching training data.

Training and validation loops in PyTorch

coreintermediate

A training loop repeats the forward/loss/backward/step cycle (§9.2) over every batch, for a number of epochs, while the model is in training mode. A validation loop runs the model on held-out data without updating weights, to check generalization (§1). A learning-rate scheduler adjusts the learning rate as training progresses. A checkpoint saves the model's state so training can resume, or the best version can be kept.

Think of it as

A real training script is the assembly of nearly every concept in the previous section into one runnable loop, plus a few additions specific to running many epochs safely. `model.train()` and `model.eval()` are not optional bookkeeping — they change real behavior in layers like dropout and batch normalization, which must behave differently during training (randomly dropping neurons, using batch statistics) versus evaluation (using all neurons, using running statistics accumulated during training) — forgetting to switch modes is a classic, silent bug. The validation loop wraps its forward passes in `torch.no_grad()` since no backward pass or weight update happens there, which also saves significant memory by skipping the autograd bookkeeping. A learning-rate scheduler changes the learning rate over the course of training — commonly decaying it as training progresses, since a large learning rate that helps escape a bad initial region can prevent fine convergence once the model is close to a good solution. A checkpoint — saving the model's `state_dict`, the optimizer's state, and the current epoch — is what makes a multi-hour or multi-day training run survivable: without it, a crash partway through means starting over from scratch, and without saving the BEST checkpoint specifically (by validation metric, not just the latest), the final saved model can be a worse epoch than one seen earlier in training.

python

best_val_loss = float("inf")
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

for epoch in range(num_epochs):
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        loss = loss_fn(model(batch.x), batch.y)
        loss.backward()
        optimizer.step()

    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for batch in val_loader:
            val_loss += loss_fn(model(batch.x), batch.y).item()

    scheduler.step()
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), "best_model.pt")   # checkpoint

What we're doing: Find the bug in a training loop that reports suspiciously perfect validation accuracy.

buggy_training_loop.pypython

for epoch in range(num_epochs):
    for batch in train_loader:
        optimizer.zero_grad()
        loss = loss_fn(model(batch.x), batch.y)
        loss.backward()
        optimizer.step()

    # validation
    val_correct = 0
    for batch in val_loader:
        predictions = model(batch.x).argmax(dim=1)
        val_correct += (predictions == batch.y).sum().item()
    print(f"epoch {epoch}: val_acc={val_correct / len(val_loader.dataset):.3f}")

8: model.train() was never called at the start of the epoch, and model.eval() was never called before validation — the model stays in whatever mode it was last in, silently.
9: No torch.no_grad() wraps the validation loop — gradients are being tracked and computational graphs built for every validation batch, wasting memory and slowing evaluation for no benefit.
12: If dropout layers are present, this bug means validation runs WITH dropout still active — hurting the reported validation accuracy rather than helping it, the opposite direction of the "suspiciously perfect" framing, which is itself a clue worth chasing down rather than assuming.

Why this works: model.train()/model.eval() and torch.no_grad() are easy to treat as boilerplate and skip, but they change what the model actually computes — a training loop that omits them can either silently under- or over-report validation performance depending on which layers the architecture uses, and either direction of that bug is genuinely hard to notice without checking for these two calls specifically.

Never calling model.eval() before the validation loop

Wrong

python

for epoch in range(num_epochs):
    for batch in train_loader:
        train_step(model, batch)
    for batch in val_loader:
        evaluate_step(model, batch)  # model still in train() mode

Better

python

for epoch in range(num_epochs):
    model.train()
    for batch in train_loader:
        train_step(model, batch)

    model.eval()
    with torch.no_grad():
        for batch in val_loader:
            evaluate_step(model, batch)

What you see: Validation metrics fluctuate more than they should batch to batch, or a model with dropout/batch norm reports validation performance that does not match the model saved and later reloaded for inference.

Why: Dropout and batch normalization are the two most common layers whose behavior depends on the model's mode — validating in the wrong mode measures a slightly different, randomized model than the one actually being trained or the one that will be deployed.

A real degree-10 model, trained past its best epoch

Training mode vs evaluation mode, what actually changes

Training mode vs evaluation mode, what actually changes
Component	model.train()	model.eval()
Dropout	Randomly zeroes neurons	All neurons active, no dropout
Batch normalization	Uses current batch's statistics	Uses running statistics accumulated during training
Gradient tracking	On (needed for backward())	Should be disabled via torch.no_grad() for memory/speed

Remember: model.train() and model.eval() change real layer behavior (dropout, batch norm), not just bookkeeping. Validation runs under torch.no_grad(). A scheduler typically decays the learning rate over training. Save the best checkpoint by validation metric, not just the last one.

GPU usage, mixed precision and memory

standardintermediate

A GPU trains a neural network far faster than a CPU because it runs many operations in parallel — but its memory is limited and shared across everything running on it. Mixed precision training uses lower-precision numbers (float16 or bfloat16) for most operations instead of float32, cutting memory use roughly in half and speeding up training on modern GPUs.

Think of it as

GPU memory, not compute, is usually the practical limit on how large a model or batch you can train — a training run does not fail because the GPU is too slow, it fails with an out-of-memory error because the model's weights, activations, gradients and optimizer state (for Adam, roughly two extra copies of every weight) all have to fit simultaneously. Mixed precision training addresses this directly: most operations run in float16/bfloat16 (half the memory and often faster on modern GPU hardware), while a few numerically sensitive operations stay in float32 to avoid precision-related instability, and a technique called loss scaling compensates for float16's narrower numeric range during backpropagation. Gradient accumulation (§9.5) is the other standard lever when a batch does not fit — smaller batches, summed gradients, less memory per step. Beyond these two techniques, memory management in practice means being deliberate about what stays on the GPU: moving data there only when needed, freeing intermediate tensors that are no longer needed, and watching for accidentally keeping a full computational graph alive (for instance, by accumulating loss values as tensors instead of calling `.item()` to pull out a plain number) long after it is needed.

python

scaler = torch.cuda.amp.GradScaler()

for batch in train_loader:
    optimizer.zero_grad()
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        loss = loss_fn(model(batch.x), batch.y)   # most ops run in float16
    scaler.scale(loss).backward()                  # loss scaling
    scaler.step(optimizer)
    scaler.update()

Remember: GPU memory is usually the real limit, not compute speed — weights, activations, gradients and optimizer state must all fit at once. Mixed precision roughly halves memory and often speeds up training; gradient accumulation is the fallback when a batch still does not fit.

Model serialization and loading

standardintermediate

Serialization saves a trained model's state to disk so it can be loaded again later, without retraining. PyTorch's recommended approach saves the model's `state_dict` — a dictionary of its learned weights — rather than the whole model object, because loading a `state_dict` back requires the same model class definition to already exist, which is exactly the safety check that prevents a mismatched architecture from loading silently wrong.

Think of it as

The core trade-off is between convenience and safety. Saving the entire model object (`torch.save(model, ...)`) is simple, but it pickles the model's class definition alongside its weights, which means loading it later requires the exact same code to be importable, and can break silently or insecurely if that code changes or the file comes from an untrusted source. Saving only the `state_dict` (a plain dictionary mapping layer names to weight tensors) is the recommended approach: you re-create the model architecture in code first, then load the weights into it — this forces the architecture to be explicit at load time, which is what catches an architecture mismatch immediately rather than producing a model that loads but behaves wrong. A checkpoint for resuming training (§10.2) typically bundles more than just the model's `state_dict` — the optimizer's `state_dict` too, so momentum and other optimizer state carry over, plus the current epoch number.

python

# saving (recommended)
torch.save(model.state_dict(), "model_weights.pt")

# loading — architecture must already be defined in code
model = MyModelClass()
model.load_state_dict(torch.load("model_weights.pt"))
model.eval()

# a full training-resume checkpoint
torch.save({
    "epoch": epoch,
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
}, "checkpoint.pt")

Remember: Save the state_dict (weights only), not the whole model object — loading it back requires the architecture to already be defined in code, which catches a mismatch immediately instead of loading something silently wrong. A resume checkpoint bundles model state, optimizer state and epoch together.

Distributed training concepts

standardintermediate

Distributed training spreads one training run across multiple GPUs, or multiple machines, because a model or dataset has grown too large — or training is too slow — for a single GPU. Data parallelism copies the whole model onto each GPU and splits the data between them. Model parallelism splits the model itself across GPUs, for models too large to fit on one.

Think of it as

The two strategies solve different bottlenecks. Data parallelism is the more common case: an identical copy of the model lives on every GPU, each processes a different slice of the batch, and after each step the gradients computed on every GPU are averaged together (an operation called all-reduce) before every copy applies the same update — the model stays in sync across GPUs because they always apply the identical, averaged gradient. This solves a speed problem (more GPUs process more data in parallel) but not a memory problem, since every GPU still needs to hold the full model. Model parallelism solves the opposite problem: when a model's weights genuinely do not fit on one GPU's memory (large language models are the common example), different layers or different parts of a layer live on different GPUs, and data flows between them as it passes through the network — this adds communication overhead between GPUs on every forward and backward pass, which data parallelism does not need since each GPU's copy is otherwise independent until the gradient sync. Most large-scale training combines both: data parallelism across groups of GPUs, model parallelism within each group.

python

# data parallelism (conceptually)
model = torch.nn.parallel.DistributedDataParallel(model)
# each GPU: forward + backward on its data slice
# then: gradients averaged across all GPUs (all-reduce)
# then: every GPU applies the identical, averaged update

Data parallelism: the real math behind "average the gradients"

Remember: Data parallelism copies the model to every GPU and splits the data — solves speed, not memory. Model parallelism splits the model itself across GPUs — solves memory for models too large for one GPU, at the cost of communication overhead. Large-scale training often combines both.

Debugging shape, device and dtype errors

standardintermediate

The three most common PyTorch errors are a shape mismatch (two tensors with incompatible dimensions for an operation), a device mismatch (one tensor on the CPU, another on the GPU), and a dtype mismatch (one tensor float32, another float16 or int64 where a float was expected). All three are caught at runtime, with an error message that usually names the exact shapes/devices/dtypes involved.

Think of it as

These three error types share a debugging approach: read the error message for the actual values involved, then trace backward to where they diverged, rather than guessing. A shape mismatch error names both shapes ('mat1 and mat2 shapes cannot be multiplied (32x10 and 5x1)') — the fix is finding which layer or reshape produced the wrong dimension, usually by printing `.shape` at each step of a forward pass until the mismatch is located. A device mismatch ('Expected all tensors to be on the same device') happens when a new tensor is created without being moved to the same device as the model and its other tensors — a batch loaded from a DataLoader lands on CPU by default and must be explicitly moved with `.to(device)` before being passed to a GPU model. A dtype mismatch commonly appears when mixing an integer tensor (often labels, which are legitimately `int64` for classification) with a float operation expecting `float32`, or when mixed precision training (§10.3) produces a `float16` tensor that then meets an operation that only accepts `float32`. All three are, deliberately, hard runtime errors rather than silent behavior — a framework that guessed how to reconcile mismatched shapes, devices or dtypes would produce numerically wrong results with no warning at all.

python

# shape debugging: print shapes at each step
print(x.shape, weight.shape)

# device fix: move every new tensor to the model's device
batch = batch.to(device)

# dtype fix: cast explicitly where needed
labels = labels.to(torch.int64)   # classification targets
x = x.to(torch.float32)            # back to full precision

A real shape-mismatch error, captured from a live interpreter

Three error types, and where to look first

Three error types, and where to look first
Error	Typical message contains	Look first at
Shape mismatch	Two specific incompatible shapes	The layer/reshape immediately before the failing operation
Device mismatch	"Expected all tensors to be on the same device"	A freshly created tensor or a DataLoader batch missing `.to(device)`
Dtype mismatch	Two specific incompatible dtypes	Labels (int64) meeting a float op, or a mixed-precision op

Remember: Shape, device and dtype mismatches are all hard runtime errors by design — read the exact values in the message, then trace backward to where they diverged, rather than guessing. Device errors are usually a missing `.to(device)`; dtype errors are often labels or mixed-precision tensors meeting the wrong operation.