Module 1 · Reading 1.8 · read before S08: PyTorch in Practice

You have now built an autograd engine, derived the optimizers, and seen why initialization and normalization matter. This section is where the course hands you the production version of all of it. PyTorch’s high-level API is small (tensors, nn.Module, torch.optim, DataLoader, and a training loop you will write so many times it becomes muscle memory), but each piece has conventions that are invisible until they bite: which mode your model is in, when gradients accumulate, what actually gets saved in a checkpoint.

This is the one systematic tour of that API the course gives. Every assignment from HW2 onward, and your semester project, assumes the vocabulary of this section. Everything runs on CPU in seconds, from synthetic data, with fixed seeds — the same discipline your homework requires.

How to read this

Tensors: NumPy with two superpowers

A torch.Tensor is an n-dimensional array, deliberately NumPy-like: shapes, dtypes, slicing, broadcasting, and reductions all transfer directly from your NumPy experience. Two additions make it the substrate of deep learning. First, a tensor can live on different devices (cpu, cuda, mps), and all operations run wherever their operands live; this is the entire GPU story from the user’s perspective. Second, a tensor can carry requires_grad=True, which makes every operation on it record onto the autograd graph you built by hand in the autograd engine; loss.backward() then fills the .grad field of every participating leaf tensor.

Two dtype habits prevent most beginner shape-and-type errors. Deep learning defaults to float32 (not NumPy’s float64): faster, half the memory, precise enough for gradient descent. And class labels for classification must be integer dtype (torch.long), because loss functions like CrossEntropyLoss use them as indices, not as numbers.

One more piece of tensor mechanics saves real confusion later: many operations return views, not copies. A tensor is a strided window onto a flat block of memory, and operations like slicing, transpose, permute, view, and squeeze often share storage. view requires a stride-compatible layout. reshape returns a view when possible and otherwise copies; .contiguous() returns the original tensor when it is already contiguous and otherwise copies into a contiguous layout. Neither operation is a blanket promise to copy. torch.from_numpy shares memory with the source array. When you require independent storage, use .clone() or .detach().clone() for a graph break as well.

PyTorch’s define-by-run design has a history worth one sentence: it descends from the Lua-based Torch7 used at NYU and Facebook, adopted the define-by-run autograd style demonstrated by Chainer (2015), and was released by Facebook AI Research in early 2017, with the NeurIPS 2019 paper by Paszke and colleagues describing the design. Its bet (that a framework should feel like ordinary Python, debuggable line by line) is why your print() statements and debugger work inside a model’s forward.

nn.Module: parameters that manage themselves

nn.Module is PyTorch’s unit of model composition, and it earns its keep through bookkeeping you no longer have to do. Assign a submodule or an nn.Parameter as an attribute in __init__, and the module registers it: model.parameters() yields every trainable tensor recursively, which is exactly what you hand to an optimizer; model.state_dict() maps names to weights for saving; model.train() / model.eval() propagate mode to every submodule. You define forward(); you call the module itself (model(x)), which routes through __call__ and its hooks; do not call model.forward(x) directly.

The first example shows the tensor basics and the registration machinery in one runnable piece:

import torch
import torch.nn as nn

torch.manual_seed(0)

# --- tensors: shapes, dtypes, broadcasting ---
x = torch.randn(4, 3)                 # 4 samples, 3 features
w = torch.randn(3, 2)                 # a 3->2 linear map
b = torch.randn(2)                    # bias, broadcast over the batch
out = x @ w + b                       # (4,3) @ (3,2) + (2,) -> (4,2)
print(f"x{tuple(x.shape)} @ w{tuple(w.shape)} + b{tuple(b.shape)} -> {tuple(out.shape)}, dtype={out.dtype}")

col = x.mean(dim=0)                   # reduce over the batch dimension
print(f"mean over dim=0 -> shape {tuple(col.shape)}")

# --- autograd on tensors: same idea as Value, elementwise ---
w.requires_grad_(True)
loss = ((x @ w) ** 2).mean()
loss.backward()
print(f"loss = {loss.item():.4f}, w.grad shape = {tuple(w.grad.shape)}")

# --- nn.Module: parameters are registered automatically ---
class TinyMLP(nn.Module):
    def __init__(self, d_in, d_hidden, d_out):
        super().__init__()
        self.fc1 = nn.Linear(d_in, d_hidden)
        self.fc2 = nn.Linear(d_hidden, d_out)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

model = TinyMLP(3, 16, 2)
total = 0
for name, p in model.named_parameters():
    print(f"{name:12s} shape={tuple(p.shape)!s:10s} requires_grad={p.requires_grad}")
    total += p.numel()
print(f"total trainable parameters: {total}")

Running this prints:

x(4, 3) @ w(3, 2) + b(2,) -> (4, 2), dtype=torch.float32
mean over dim=0 -> shape (3,)
loss = 0.5846, w.grad shape = (3, 2)
fc1.weight   shape=(16, 3)    requires_grad=True
fc1.bias     shape=(16,)      requires_grad=True
fc2.weight   shape=(2, 16)    requires_grad=True
fc2.bias     shape=(2,)       requires_grad=True
total trainable parameters: 98

Note the conventions on display: the batch dimension comes first everywhere; nn.Linear(d_in, d_out) stores its weight as (d_out, d_in) and computes \(x W^\top + b\); .item() extracts a Python float from a scalar tensor (use it for logging; holding on to the tensor itself keeps the whole autograd graph alive); and every registered parameter defaults to requires_grad=True. When you need a quick stack of layers with no custom logic, nn.Sequential builds a module from a list; the moment you need branching, multiple inputs, or anything nonlinear in structure, write the nn.Module subclass.

People. PyTorch began as an intern project: Adam Paszke started the effort in 2016 while an undergraduate intern working with Soumith Chintala at Facebook AI Research, building on the existing Torch ecosystem; Sam Gross and Gregory Chanan were among the other core early developers. The NeurIPS 2019 paper (“PyTorch: An Imperative Style, High-Performance Deep Learning Library”) is the standard citation and a readable account of the design choices — chiefly, that researcher ergonomics were worth reclaiming from compiler-style frameworks.

Data: Dataset and DataLoader

Training needs data in shuffled minibatches, and PyTorch splits that job in two. A Dataset answers two questions (len(self) and self[i]) — that is, how many examples exist and what example i looks like. A DataLoader wraps any dataset and handles batching, shuffling, and (on real projects) parallel loading with worker processes. For tensors already sitting in memory, TensorDataset adapts them to the Dataset interface with no code.

The division matters because each half scales independently: your dataset can be “read image from disk and augment it” without knowing anything about batching, and your training loop consumes (xb, yb) pairs without knowing where they came from. Two knobs matter immediately: shuffle=True for training (a fresh random order each epoch, which the stochastic-gradient analysis in optimizers quietly assumes) and shuffle=False for evaluation, where order is irrelevant and reproducibility is nice.

On real projects, the data pipeline is also where training runs quietly lose most of their speed, so the performance knobs deserve a conceptual map even though every example in this section is too small to need them. The disease is starvation: a GPU that finishes its forward-backward pass and then idles, waiting for the next batch to be read from disk, decoded, and augmented in Python. The first remedy is num_workers: set it above zero and the DataLoader spawns that many worker processes, each running your dataset’s __getitem__ in parallel and handing finished batches to the training process; loading overlaps with computation instead of alternating with it. The second is pin_memory=True, which stages batches in page-locked (“pinned”) host memory; copies from pinned memory to the GPU are faster and can run asynchronously with compute (pass non_blocking=True to the .to(device) call to get the overlap). Two companions round out the set: persistent_workers=True keeps the worker processes alive between epochs instead of re-spawning them, and prefetch_factor controls how many batches each worker keeps ready in advance.

The map matters more than the numbers, and it cuts both ways. Workers pay off exactly when __getitem__ does real work: disk reads, image decoding, augmentation. For a dataset like this section’s, whose tensors already sit in memory, worker processes add serialization overhead for zero benefit, and num_workers=0 is genuinely fastest; a recurring mistake is to cargo-cult num_workers=8 onto a small assignment and find it slower. The honest procedure is the one this course keeps recommending for optimization of any kind: measure first (does GPU utilization sag between steps? does epoch time change when you double the workers?), then tune the knob the measurement indicts.

The canonical training loop

Every PyTorch training script contains the same seven-line core (whose four-line heart is the zero_grad/forward/backward/step cycle), and you should be able to write it from memory:

for each epoch:
    model.train()
    for xb, yb in train_loader:
        opt.zero_grad()          # 1. clear accumulated gradients
        loss = loss_fn(model(xb), yb)   # 2. forward pass
        loss.backward()          # 3. backward pass: fill .grad
        opt.step()               # 4. update parameters

Each line is something you have already built. zero_grad() exists because gradients accumulate by design — your Value engine did the same with +=. backward() is the topological-order sweep from the autograd engine. step() is one update of whichever rule from optimizers you chose. The only new discipline is modes: for ordinary metric evaluation, use both model.eval() (switches dropout and BatchNorm behavior) and torch.no_grad() or torch.inference_mode() (stops graph recording). They are independent. Evaluation-like analyses that differentiate with respect to inputs or parameters (saliency maps and adversarial attacks, for example) still use model.eval() but must leave gradients enabled.

Here is the complete program: synthetic two-spiral data, a model with dropout (so modes matter), the loop, evaluation, and checkpointing. It runs in about two seconds on CPU.

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

torch.manual_seed(0)

# --- 1. Synthetic data: two interleaved spirals, 2 classes ---
def make_spirals(n_per_class, noise=0.08):
    t = torch.linspace(0.5, 3 * torch.pi, n_per_class)
    xs, ys = [], []
    for label in (0, 1):
        r = t / (3 * torch.pi) * 2.5
        angle = t + label * torch.pi          # second spiral rotated 180 degrees
        pts = torch.stack([r * torch.cos(angle), r * torch.sin(angle)], dim=1)
        xs.append(pts + noise * torch.randn(n_per_class, 2))
        ys.append(torch.full((n_per_class,), label, dtype=torch.long))
    return torch.cat(xs), torch.cat(ys)

X, y = make_spirals(400)
perm = torch.randperm(len(X))                  # shuffle before splitting
X, y = X[perm], y[perm]
n_train = int(0.8 * len(X))
train_ds = TensorDataset(X[:n_train], y[:n_train])
val_ds = TensorDataset(X[n_train:], y[n_train:])
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=256)

# --- 2. Model ---
class SpiralNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 64), nn.ReLU(), nn.Dropout(0.1),
            nn.Linear(64, 64), nn.ReLU(), nn.Dropout(0.1),
            nn.Linear(64, 2),
        )

    def forward(self, x):
        return self.net(x)          # raw logits; CrossEntropyLoss applies softmax

model = SpiralNet()
loss_fn = nn.CrossEntropyLoss()
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)

# --- 3. The canonical loop ---
def evaluate():
    model.eval()                    # dropout off, etc.
    correct = n = 0
    total_loss = 0.0
    with torch.no_grad():           # no graph needed for evaluation
        for xb, yb in val_loader:
            logits = model(xb)
            total_loss += loss_fn(logits, yb).item() * len(xb)
            correct += (logits.argmax(dim=1) == yb).sum().item()
            n += len(xb)
    return total_loss / n, correct / n

for epoch in range(1, 151):
    model.train()                   # dropout on
    running = 0.0
    for xb, yb in train_loader:
        opt.zero_grad()             # 1. clear old gradients
        loss = loss_fn(model(xb), yb)   # 2. forward
        loss.backward()             # 3. backward
        opt.step()                  # 4. update
        running += loss.item() * len(xb)
    if epoch % 30 == 0:
        val_loss, val_acc = evaluate()
        print(f"epoch {epoch:3d}: train loss {running / n_train:.4f} | "
              f"val loss {val_loss:.4f} | val acc {val_acc:.3f}")

# --- 4. Save, reload into a fresh model, verify ---
ckpt = "spiralnet.pt"
torch.save(model.state_dict(), ckpt)

model = SpiralNet()                 # fresh, randomly initialized model
_, acc_fresh = evaluate()
model.load_state_dict(torch.load(ckpt))
_, acc_loaded = evaluate()
print(f"fresh model val acc: {acc_fresh:.3f} | after load_state_dict: {acc_loaded:.3f}")

Running this prints:

epoch  30: train loss 0.2941 | val loss 0.2381 | val acc 0.956
epoch  60: train loss 0.1175 | val loss 0.0597 | val acc 1.000
epoch  90: train loss 0.0586 | val loss 0.0265 | val acc 1.000
epoch 120: train loss 0.0549 | val loss 0.0175 | val acc 1.000
epoch 150: train loss 0.0502 | val loss 0.0120 | val acc 1.000
fresh model val acc: 0.481 | after load_state_dict: 1.000

Walk through what the numbers show. For this seed and split, the MLP classifies all 160 validation examples correctly from epoch 60 onward. That is evidence about this run, not a universal optimizer or architecture ranking. Validation loss keeps falling after accuracy reaches 1.000 because confidence can still change without changing the predicted class. Also, training loss is measured with dropout active while evaluate switches dropout off, so the two loss values are not computed under identical model behavior. Note that the model outputs raw logits: nn.CrossEntropyLoss fuses the softmax and the log-likelihood internally for numerical stability, a design covered in losses and gradient descent; putting a Softmax layer at the end of your model and then applying CrossEntropyLoss is a classic silent error.

Pitfall. loss.item() versus loss: logging the tensor (losses.append(loss)) instead of the float (loss.item()) keeps every iteration’s entire computation graph reachable, and memory climbs until the process dies — a leak that surfaces only on long runs. The same trap applies to any tensor you store for later: call .item() for scalars or .detach() for tensors before stashing them.

Saving and loading: state_dict, not the model

The last four lines of the example are the checkpointing idiom, and the verification is the point: a freshly constructed SpiralNet scores 0.481 (chance level, random weights), and after load_state_dict it scores exactly the trained 1.000. The recommended pattern is to save the state_dict (a plain name-to-tensor dictionary) rather than pickling the whole model object, because the dictionary is portable across code refactors: loading only requires that names and shapes match the current class definition. Recreate the architecture in code, then pour the weights in.

It is worth knowing that state_dict holds more than parameters() does. Modules can also register buffers (tensors that belong to the model’s state but receive no gradients) via register_buffer, and the canonical example is one you met in regularization: BatchNorm’s running mean and variance. Those statistics are learned from data (accumulated during training, used at eval time), so they must ship with the checkpoint, but no optimizer ever updates them, so they are not parameters. state_dict includes both kinds; model.parameters() yields only the trainable ones. This distinction explains an otherwise-mystifying failure: a BatchNorm model whose parameters were copied by hand (rather than through load_state_dict) evaluates badly, because the running statistics stayed at their defaults — the weights arrived, the buffers did not.

Three practical notes complete the picture. If you want to resume training rather than just deploy, save the optimizer’s state too (opt.state_dict()): Adam’s moment estimates from optimizers are state, and restarting without them changes trajectories. Since PyTorch 2.6, torch.load defaults to weights_only=True, refusing to unpickle arbitrary objects. The threat is concrete: unpickling can execute code the file’s author chose, so loading an untrusted full-model checkpoint means running an untrusted program, while a plain tensor dictionary contains nothing executable — another reason it is the right thing to save. And after loading a model for inference, call model.eval() before you trust its outputs; load_state_dict restores weights, not modes.

Historical note. The dominant framework when PyTorch appeared was TensorFlow (released by Google in 2015), which then required building a static graph and executing it inside a session — powerful for deployment, painful for debugging. PyTorch’s 2017 release made define-by-run the default experience for researchers, and by most accounts it became the dominant framework in ML research over the following years; TensorFlow 2 (2019) itself moved to eager execution by default. The episode is a standing lesson that developer experience is a research accelerant, not a luxury.

Reading the error message: shapes, dtypes, devices

Most of the time you will spend “debugging PyTorch” is really spent debugging three properties of tensors: shape, dtype, and device. The errors they raise look intimidating and are actually among the most informative in the Python ecosystem, once you know that the numbers in the message are the answer, not noise. The program below triggers the two loudest failures on purpose, then demonstrates the quiet one that raises no error at all:

import warnings
import torch
import torch.nn as nn

torch.manual_seed(0)

# 1. Shape mismatch: the error names both shapes -- read them.
layer = nn.Linear(10, 3)
x = torch.randn(4, 8)                       # 8 features, layer expects 10
try:
    layer(x)
except RuntimeError as e:
    print(f"[shape] {e}")

# 2. Wrong label dtype at the loss.
logits = torch.randn(4, 3)
labels_float = torch.tensor([0.0, 1.0, 2.0, 1.0])   # should be torch.long
try:
    nn.CrossEntropyLoss()(logits, labels_float)
except RuntimeError as e:
    print(f"[dtype] labels={labels_float.dtype}; expected=torch.long class indices")

# 3. The silent one: broadcasting turns (8,) vs (8,1) into an (8,8) comparison.
pred = torch.randn(8)
target = torch.randn(8, 1)
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    wrong = nn.MSELoss()(pred, target)      # no error -- but not what you meant
right = nn.MSELoss()(pred, target.squeeze(1))
print(f"[broadcast] warning issued: {'target size' in str(caught[0].message)}")
print(f"[broadcast] loss with (8,) vs (8,1): {wrong.item():.4f}   after squeeze: {right.item():.4f}")
print(f"[broadcast] (pred - target).shape = {tuple((pred - target).shape)}  <- 64 pairwise differences")

Running this prints:

[shape] mat1 and mat2 shapes cannot be multiplied (4x8 and 10x3)
[dtype] labels=torch.float32; expected=torch.long class indices
[broadcast] warning issued: True
[broadcast] loss with (8,) vs (8,1): 3.5749   after squeeze: 5.0210
[broadcast] (pred - target).shape = (8, 8)  <- 64 pairwise differences

Take the three in order. The shape error (4x8 and 10x3) gives you both operands: 4x8 is your batch (4 samples, 8 features) and 10x3 is the layer’s weight seen as a 10 -> 3 map, so the mismatch (8 versus 10) is visible in the message itself, and the only question left is which of the two numbers is wrong. The habit to build: do not stare at code, print shapes. Sprinkling print(x.shape) at the boundaries of a model (or running one dummy batch through each submodule right after constructing it) resolves nearly every shape bug in minutes, and the batch-dimension-first convention means you can usually spot the offending axis on sight.

The dtype error is the one from the tensors section: classification targets must be torch.long because the loss uses them to index rows of the logits. The message names the expectation and the offense. Dtype bugs cluster at two borders: labels (float instead of long, often from reading a CSV with NumPy defaults) and NumPy interop (np.float64 arriving where float32 is expected; torch.from_numpy preserves the 64-bit dtype and the mismatch surfaces later, in a matmul).

The third case is the dangerous one, because the program runs. pred has shape (8,), target has shape (8, 1); subtracting them broadcasts both to (8, 8), so the “loss” averages all 64 pairwise differences instead of 8 matched ones — a different number, differently wrong every batch, and gradient descent will happily optimize it. PyTorch does emit a UserWarning naming the two sizes (the code above captures it), but warnings scroll past in training logs. The defenses are mechanical: squeeze/unsqueeze deliberately at data boundaries, and an assert pred.shape == target.shape in your loss-adjacent code, one line that would have converted this silent bug into a loud one.

The third property, device, cannot be demonstrated on the CPU-only machines these readings target, but the error reads like the others: Expected all tensors to be on the same device, but found at least two devices, followed by the two device names. It appears the moment you move a model to the GPU and forget the data (or vice versa), and the fix is the standard idiom: create a single device variable at the top of the script, move the model once (model.to(device)), and move each batch inside the loop (xb.to(device)). One subtlety worth knowing before you first touch a GPU in HW2: Module.to(device) moves the module’s parameters in place, but Tensor.to(device) returns a new tensor — writing xb.to(device) without assigning the result is a no-op, and the error message will keep insisting your data is on the CPU because it is.

Pitfall. When a shape bug produces no error anywhere near its cause, binary-search with shapes, not with intuition: print x.shape after each block of the model, compare against what you expect, and the first line where they diverge is the bug. Resist the urge to “fix” a shape error by adding .view(...) or .reshape(...) until the message goes away — reshaping to silence an error, without knowing which axis meant what, is how a (batch, features) tensor becomes a scrambled (features, batch) one that trains to garbage with no error at all.

Reproducibility: seeds, determinism, and their limits

Every example in this section starts with torch.manual_seed(0), and it is worth being precise about what that buys. PyTorch keeps a global random number generator (per device) that feeds every stochastic operation that is not given its own generator: parameter initialization inside nn.Linear, torch.randn for synthetic data, DataLoader shuffling, dropout masks. Seeding it pins the entire sequence, so a run becomes a deterministic function of the code and the seed. The program below trains a small classifier twice with the same seed and once with a different one:

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

def train_once(seed):
    torch.manual_seed(seed)                  # governs init AND DataLoader shuffling
    torch.use_deterministic_algorithms(True) # error out on nondeterministic kernels
    gen = torch.Generator().manual_seed(0)   # data itself fixed across runs
    X = torch.randn(256, 4, generator=gen)
    y = (X.sum(dim=1) > 0).long()
    loader = DataLoader(TensorDataset(X, y), batch_size=32, shuffle=True)
    model = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 2))
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()
    for _ in range(5):
        for xb, yb in loader:
            opt.zero_grad()
            loss = loss_fn(model(xb), yb)
            loss.backward()
            opt.step()
    return loss.item()

a, b, c = train_once(0), train_once(0), train_once(1)
print(f"seed 0, run 1: final batch loss = {a:.10f}")
print(f"seed 0, run 2: final batch loss = {b:.10f}")
print(f"seed 1:        final batch loss = {c:.10f}")
print(f"run 1 == run 2: {a == b}   (bitwise identical)")

Running this prints:

seed 0, run 1: final batch loss = 0.6037580371
seed 0, run 2: final batch loss = 0.6037580371
seed 1:        final batch loss = 0.5568858385
run 1 == run 2: True   (bitwise identical)

Two runs, hundreds of random events each (initialization of four parameter tensors, five epochs of shuffling, every gradient step), and the final losses agree to the last bit. The third run shows the flip side: change nothing but the seed and the loss lands somewhere else entirely. Both facts matter for your homework and project. Identical-seed reproducibility is what makes debugging possible (rerun, get the same failure) and what any step-for-step comparison against a PyTorch reference relies on. Different-seed variation is what makes single-run comparisons meaningless: if seed alone moves your metric by more than your proposed improvement does, you have measured noise — the discipline of averaging over seeds returns in experiments and reproducibility.

Note the separation trick in the code: the data is drawn from its own explicitly seeded torch.Generator, so the dataset is identical across all three runs and only initialization and shuffling vary with the seed. Passing dedicated generators to the pieces you want pinned (a DataLoader accepts one too) is cleaner than trying to reason about the global stream, in which every added randn shifts everything after it.

One corner of seeding deserves a warning: randomness inside DataLoader workers. PyTorch assigns each worker a distinct torch seed. For NumPy, Python’s random, or another library used inside __getitem__, follow the current DataLoader/reproducibility documentation: initialize that library from torch.initial_seed() in worker_init_fn when you need explicit, version-stable control. Multiprocessing start methods and library versions affect inherited external RNG state, so do not assume that seeding torch alone specifies every random stream.

The other line, torch.use_deterministic_algorithms(True), addresses a different threat. On CPU with ordinary layers, seeded runs are deterministic anyway. On GPUs, some kernels are not: floating-point addition is not associative ((a + b) + c and a + (b + c) can differ in the last bit), so any kernel whose parallel reduction order varies between launches (some scatter and pooling backward passes, notably) can return different results for identical inputs. The flag tells PyTorch to use deterministic implementations where they exist and to raise an error naming the offending operation where they do not, letting you decide whether to accept the slower path (some also need the CUBLAS_WORKSPACE_CONFIG environment variable set, as the error message will tell you). Full GPU determinism also wants torch.backends.cudnn.benchmark = False, since the benchmark mode auto-selects convolution algorithms by timing them, and the winner can differ run to run. Two honest caveats close the topic. First, determinism costs speed, and the cost has no single ballpark: the flag swaps only the kernels your model actually uses, so the slowdown ranges from nothing (this section’s CPU examples already run deterministic kernels) to substantial (a kernel whose only deterministic implementation is a slow fallback), and disabling benchmark mode costs whatever its autotuning was winning. The practical move is to time one epoch each way on your own model; the common pattern keeps determinism on while debugging or comparing runs, and off for long production training where throughput dominates. Second, no flag makes results reproducible across different hardware, driver, or PyTorch versions — bitwise reproducibility is a claim about one environment, which is why serious projects record the environment alongside the seed.

Going deeper

torch.compile, or the graph strikes back. The autograd engine section told the history: static graphs lost to define-by-run on ergonomics, then compilers came partway back. PyTorch 2.0 (released in 2023) is that return trip: model = torch.compile(model) is the entire user interface. Underneath, a component called TorchDynamo intercepts your Python bytecode as it runs, extracts the tensor operations into a graph, and hands it to a backend called TorchInductor, which generates fused kernels (Triton for GPUs, C++ for CPUs); wherever capture would change semantics, it falls back to ordinary eager execution for that fragment, so correctness does not depend on your code being capturable. The win comes mostly from fusion: eager PyTorch runs one kernel per operation, each reading and writing memory, while a compiler can weld a chain of elementwise operations into one kernel that touches memory once — the PyTorch 2 paper (Ansel et al., ASPLOS 2024) reports geometric-mean training speedups of about 1.4x on an A100 GPU across 180+ real models. The costs: the first iterations are slow (compilation happens then), changing input shapes can trigger recompilation, and stack traces get harder to read, which is why the practical advice is to develop and debug in eager mode, then add the one line when the model is stable. Nothing in your mental model changes: autograd, modules, and the training loop all work identically under compilation.

Mixed precision. Modern accelerators run half-precision matrix math several times faster than float32, and modern training exploits this with automatic mixed precision (AMP): under torch.autocast, each operation runs in the lowest precision that is safe for it — matmuls in float16 or bfloat16, loss computations and reductions in float32. With classic float16 the small exponent range means gradients can underflow to zero, so AMP pairs autocast with a GradScaler that multiplies the loss by a large factor before backward() and divides it back out before step(); bfloat16 (float32’s exponent range with less mantissa) usually needs no scaler, which is why it has become the default on hardware that supports it. AMP is standard practice for any serious GPU training and returns in efficiency and scale.

More than one GPU. The standard first step beyond a single device is distributed data parallelism (DistributedDataParallel, DDP): one process per GPU, each holding a full copy of the model and consuming a different shard of each batch; after every backward pass the processes average their gradients (an all-reduce collective), so every copy takes the same optimizer step and the whole ensemble behaves like one model trained on the combined batch, which is exactly the large-batch regime whose optimizer consequences (linear scaling, warmup) the optimizers section discussed. Models too large for one GPU’s memory need the further step of sharding parameters and optimizer state across devices (fully sharded data parallelism, FSDP); Chapter 4 picks this up. For this course’s homework, everything fits on one device — but the training loop you learned here is unchanged in all of these settings, which is much of why it is worth committing to memory.

Common pitfalls

  • Symptom: evaluation accuracy is noisy across identical runs, or worse than training accuracy by a suspicious constant margin. Cause: model left in training mode — dropout active during evaluation. Fix: model.eval() inside your evaluation function (and model.train() at the top of each epoch, since evaluation flipped the mode).
  • Symptom: loss decreases but 10–100x slower than expected. Cause: missing opt.zero_grad(), so each step applies the sum of all gradients so far — or a Softmax layer feeding CrossEntropyLoss, squashing gradients. Fix: follow the four-line core exactly; the loss function wants raw logits.
  • Symptom: RuntimeError: expected scalar type Long but found Float (or similar dtype complaints) at the loss. Cause: classification targets stored as floats. Fix: labels must be torch.long; create them with dtype=torch.long as make_spirals does.
  • Symptom: memory usage grows steadily during training or evaluation. Cause: retaining graph-connected loss/output tensors across iterations, using retain_graph=True, or recording evaluation graphs unnecessarily. After a normal backward pass PyTorch frees most saved intermediates, but retained tensors and explicit graph retention can still keep graph structure or state alive. Fix: use .item() for logged scalars, detach stored predictions, and disable gradients for ordinary inference.
  • Symptom: a reloaded checkpoint gives chance-level predictions. Cause: weights never actually loaded (result of SpiralNet() used directly, load_state_dict return value ignored after loading into the wrong object) or architecture changed since saving so names silently mismatched with strict=False. Fix: verify after loading, as the example does — one known input, or a held-out metric, compared against the pre-save value.

Try it yourself

  1. Remove model.eval() from evaluate() in the end-to-end example and rerun. How much does reported validation accuracy change, and why does it now vary between calls on identical data?
  2. Extend the checkpoint to a dictionary containing model.state_dict(), opt.state_dict(), and the epoch number. Stop training at epoch 75, restore everything, continue to 150, and compare the final validation loss to the uninterrupted run.
  3. Replace TensorDataset with your own Dataset subclass implementing __len__ and __getitem__ that generates each spiral point on the fly from its index. Verify the training curve is unchanged when seeds are controlled.
  4. Add a cosine schedule with warmup (from the optimizers section) to the end-to-end example and log sched.get_last_lr() at each printout. Does it change where the model reaches val acc 1.000?
  5. Run the spiral training five times with seeds 0–4 (seeding both the data split and the model) and report the mean and range of final validation accuracy. Then introduce the (8,)-versus-(8,1) broadcasting bug from the error-message section into a regression variant of the loop and confirm it trains without any error while producing a wrong loss.

Further reading

Before class you should be able to

  • Write the four-line training-loop core from memory and connect each line to what you built in the S05 and S06 readings.
  • Explain why model.eval() and disabled gradient recording are independent, why ordinary metric evaluation uses both, and when an evaluation-time analysis still needs gradients.
  • Explain why models output raw logits to nn.CrossEntropyLoss, and why a trailing Softmax layer is a classic silent error.
  • Save and reload a state_dict, explain buffers versus parameters, and verify that the load actually happened.
  • Read shape and dtype error messages as answers, and spot the silent (8,)-versus-(8, 1) broadcasting bug.
  • Explain loss.item() versus storing the tensor, and the memory leak the latter causes.
  • Say what torch.manual_seed pins, what identical seeds guarantee, and why comparisons need multiple seeds.

Connects to

  • This session is covered by Quiz 1 (Module 1: S01–S08), given at S09.
  • HW1 (micro-autograd) is due today. From here on every assignment is written against this API, so the nn.Module, state_dict, and DataLoader conventions in this reading become immediately load-bearing.
  • The loop learned here is the skeleton for HW2–HW4 and the semester project; debugging it at scale is the subject of the training-practice session (S12).

The PDF posted in Canvas is the assigned version of this reading; it is built from this page.

← S07 · Regularization, Initialization, Normalization · All readings · S09 · Convolution →