Tensors that know their shape, dtype and device. Autograd that writes backward() for you. nn.Module trees. And five lines of training loop that trained GPT-4, Stable Diffusion and LLaMA — every concept one you already built by hand, now running on C++/CUDA kernels.
Shape is how the numbers are indexed; dtype is how precisely each one is stored (float32 = 4 bytes, float16 = 2, int8 = 1); device is where the arithmetic happens. Every operation requires all its tensors on one device and in compatible shapes — which is why print(x.shape, x.dtype, x.device) answers half of all PyTorch questions.
shape (64, 784) · float32 · cuda:0Every operation on a tensor that requires grad appends one entry to a tape. Calling .backward() on the loss walks the tape in reverse, multiplying each local derivative into a running gradient: x² + 3x summed gives x.grad = 2x + 3 = [5, 7, 9] with no hand-written backward() anywhere. Leaves accumulate .grad, so zero it before the next pass.
z = Σ(x² + 3x) = 32 → x.grad = [5, 7, 9]optimizer.zero_grad(), forward, loss, loss.backward(), optimizer.step() — in that order, for every batch, for every epoch. Dropout and batch norm need model.train() while training and model.eval() with torch.no_grad() while testing. The architecture changes, the data changes, these five lines do not: they trained GPT-4, Stable Diffusion and LLaMA.
zero → forward → loss → backward → stepPyTorch is your Lesson 10 framework with the parts you never wrote: tensors that record their own arithmetic, kernels written in C++/CUDA, and a state_dict that outlives the process — wrapped around the exact Module, forward, parameters, backward and step interface you designed yourself.
By the end you will be able to read any tensor’s shape, dtype and device and predict the error before it happens; compute a small backward pass by hand (x.grad = 2x + 3 = [5, 7, 9]) and check it with a finite difference; count the parameters of any Linear stack — Linear(4, 3) = 15, the source’s MNIST MLP = 235,146; write the canonical five-line loop from memory and say what each line does; hand-compute cross-entropy on raw logits and explain why softmax before it is a silent bug; move a model and its batches onto one device; save and load a state_dict; and translate every component of the mini framework into its PyTorch equivalent.
In Lesson 10 you assembled Linear layers, ReLU, dropout, Adam and a training loop out of pure Python. It works. It is also about 500× slower than the same model in PyTorch — and the gap is not in the math, it is in how many times you make Python do the work.
Your mini framework processes one sample at a time with nested Python loops. Every multiply-add passes through the interpreter. PyTorch dispatches the same operations to optimized C++/CUDA kernels that chew through a whole batch at once, and on a GPU run thousands of arithmetic lanes in parallel. The source’s own measurements on the same MNIST MLP: ~300 s/epoch for the mini framework versus ~0.5 s/epoch for PyTorch. On a single NVIDIA A100, PyTorch trains ResNet-50 (25.6M parameters) on ImageNet (1.28M images) in about 6 hours; the pure-Python version would need roughly 3,000 hours — if it did not run out of memory first.
Speed is not the only gap. Your framework has no GPU support. No automatic differentiation — you hand-wrote backward() for every module. No serialization, no distributed training, no mixed precision, no way to debug gradient flow without print statements. PyTorch fills every one of those gaps while keeping the exact mental model you already built: Module, forward(), parameters(), backward(), optimizer.step(). The concepts transfer one-to-one. The syntax is nearly identical. The difference is that PyTorch wraps a decade of systems engineering behind the same interface you designed from scratch.
Why did PyTorch win? In 2015 TensorFlow made you define a static computation graph before running anything: build the graph, compile it, then feed data through. Debugging meant staring at graph visualizations, and changing the architecture meant rebuilding the graph. PyTorch launched in 2017 with eager execution: you write Python, it runs immediately. y = model(x) actually computes y right now. That meant print() worked, pdb worked, an if statement inside a forward pass worked. By 2020 the market had decided: PyTorch’s share of ML research papers went from 7% in 2017 to over 75% in 2022, and TensorFlow 2.x adopted eager execution in response. The lesson generalises: developer experience compounds — a framework that is 10% slower but 50% faster to debug wins every time.
| Mini framework (Lesson 10) | PyTorch |
|---|---|
model = Sequential(Linear(784, 256), ReLU(), …) | model = nn.Sequential(nn.Linear(784, 256), nn.ReLU(), …) |
pred = model.forward(x) | pred = model(x) |
optimizer.zero_grad() | optimizer.zero_grad() |
grad = criterion.backward(); model.backward(grad) | loss.backward() |
optimizer.step() | optimizer.step() |
no GPU support | model.to("cuda") |
manual backward() for every module | autograd handles everything |
Same multiply-adds, different number of trips through the interpreter. One tick is one dispatch: the mini framework calls Python batch-size times; PyTorch launches a single kernel that does the whole batch at once.
The gap is not one clever trick. It is fewer dispatches, vectorized C++/CUDA kernels, and — on a GPU — thousands of arithmetic lanes running at once.
The takeaway. You are not learning a new subject. You are learning the production implementation of the subject you already know. Every chapter from here on maps back to something you built by hand — and the mapping is close enough that you can read PyTorch code like Python, because it is Python.
Every PyTorch value is a multi-dimensional array that carries three facts with it: its shape, its dtype and its device. Learn to read those three at a glance and most PyTorch error messages start explaining themselves.
The creation calls look like the arrays you know from NumPy, and that is not a coincidence — but PyTorch tensors add one thing NumPy arrays do not have: every operation is recorded for autograd, and every value knows which device it lives on.
import torch
x = torch.zeros(3, 4) # shape (3, 4), dtype float32, device cpu
x = torch.randn(2, 3, 224, 224) # batch of 2 RGB images, 224×224
x = torch.tensor([1, 2, 3]) # from a Python list
x = torch.randn(2, 3, 4)
x.view(2, 12) # reshape to (2, 12) — must be contiguous
x.reshape(6, 4) # reshape to (6, 4) — works always (copies if needed)
x.permute(2, 0, 1) # reorder dimensions → (4, 2, 3)
x.unsqueeze(0) # add a size-1 dimension → (1, 2, 3, 4)
x.squeeze() # remove size-1 dimensions
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.randn(3, 4, device=device)
x = x.to("cuda") # move one tensor
x = x.cpu()
model = model.to(device) # recursively move every parameter and bufferShape is how the numbers are indexed. A scalar is (), a vector is (n,), a matrix is (m, n), and a batch of images is (batch, channels, height, width). The batch dimension comes first by convention, and MNIST’s 28×28 images arrive flattened as 784 = 28 × 28 features, so a training batch is (64, 784).
Dtype controls precision, memory and what can be computed. The four you will actually meet:
| dtype | bits | precision | use case |
|---|---|---|---|
float32 | 32 | ~7 decimal digits | default training |
float16 | 16 | ~3.3 decimal digits | mixed precision |
bfloat16 | 16 | float32 range, less precision | LLM training |
int8 | 8 | −128 … 127 | quantized inference |
The dtype choice is a memory decision, and the numbers are easy to check. The source’s MNIST model has 235,146 parameters; at 4 bytes each that is 940,584 bytes ≈ 918.5 KiB; in float16 it is 459.3 KiB; stored as int8 for inference it would be 229.6 KiB. A single training batch of 64 flattened images is 64 × 784 × 4 = 200,704 bytes ≈ 196 KiB, while the whole 60,000-image training set is 179.4 MiB — which is why the DataLoader streams batches after the first epoch instead of holding everything in the accelerator’s memory.
Device is where the arithmetic happens. There is no silent fallback: every operation requires all its tensors on the same device. Move the model once (model = model.to(device)), then move every batch inside the training loop. This is the #1 PyTorch error beginners hit, and its message is worth recognising on sight:
Reshaping is free. Operations like view, reshape, unsqueeze and squeeze change only the metadata — the 24 numbers of a (2, 3, 4) tensor never move. permute is the odd one out: it changes which element sits at each logical position, which is why it has its own chapter of consequences below.
Contiguity is not a vibe; it is a striding statement. Row-major strides for shape (2, 3, 4) are (12, 4, 1): element (a, b, c) lives at storage offset a·12 + b·4 + c·1.
y.reshape(2, 12) succeeds because it is allowed to copy: it materialises the 24 values in logical order into a fresh contiguous block, then re-chunks that. For 24 numbers the copy is invisible; for a 4K image batch it is real work, which is why view is preferred when you know the tensor is contiguous. PyTorch even has a check for it: y.is_contiguous() is False after the permute.
Start from the source’s torch.randn(2, 3, 4). Reshape, permute and squeeze it, change its precision, and move it between devices — the storage strip at the bottom never moves.
Every operation needs all tensors on the same device. This is the #1 beginner error in PyTorch: Expected all tensors to be on the same device.
You call y = x.permute(2, 0, 1) and then y.view(2, 12). PyTorch raises 'view size is not compatible with input tensor's size and stride'. What happened, and what should you use instead?
Every operation on a tensor that requires gradients appends one entry to a tape. Calling .backward() replays the tape in reverse, multiplying local derivatives along every edge — the chain rule, automated for any graph you can write in Python.
In the mini framework you implemented a backward() method for every module: the Linear layer knew its own gradient formula, ReLU knew its mask, and you wired them together by hand. PyTorch removes that job entirely. During the forward pass it records a directed acyclic graph — the computational graph, or “tape” — with one node per operation and one edge per dependency. .backward() starts at the loss with gradient 1 and walks the graph in reverse, multiplying each stored local derivative into the running gradient.
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2 + 3 * x
z = y.sum()
z.backward()
print(x.grad) # dz/dx = 2x + 3 → tensor([5., 7., 9.])
w = torch.randn(3, requires_grad=True)
for step in range(3):
loss = (w ** 2).sum() # forward: builds a fresh tape
loss.backward() # backward: replays it in reverse
with torch.no_grad(): # updates are not part of the graph
w -= 0.1 * w.grad
w.grad.zero_() # gradients accumulate — clear themThe tape explains the three rules that trip up every beginner:
1 · Only leaf tensors keep gradients. A leaf is a tensor you created (not one produced by an operation), and only leaves with requires_grad=True accumulate into .grad. Intermediate results such as y above are part of the tape but are not asked to store gradients — unless you call y.retain_grad().
2 · Gradients accumulate by default. .backward() adds to whatever is already in .grad; it does not overwrite. That is why every training loop starts with optimizer.zero_grad(). It is also the mechanism behind gradient accumulation, the trick that lets small GPUs simulate large batches.
3 · torch.no_grad() switches the tape off. During evaluation you do not need gradients, and recording them costs memory proportional to the activations. Wrap the whole evaluation loop in with torch.no_grad(): and the tape is never built.
Every local derivative above is exactly the kind of formula you wrote by hand in Lesson 03. Autograd’s contribution is bookkeeping: it stores each one as it goes, then multiplies them along the tape. For a network with millions of parameters, the bookkeeping is the entire problem — and it is now free.
Worked example B — what a training step does. The source’s second demo minimises loss = Σw² for w = [1, −2, 0.5], which is the whole optimisation story in miniature. The gradient of a sum of squares is 2w — one line, checkable by eye:
One backward pass produced all three gradients at once; one with torch.no_grad(): block applied the update without recording the update itself on the tape; one .zero_() cleared the slate for the next pass. You have now seen the entire training loop — forward, backward, step, zero — in six tensor values.
Step through the source’s own example — x² + 3x summed — and watch the reverse pass accumulate each local derivative into x.grad.
This is tape-based autodiff: every operation appended one entry on the forward pass; .backward() walks it in reverse. Your mini framework wrote these steps with a pencil — PyTorch wrote them in C++.
You can watch accumulation inside a single training pass too: the four micro-batch gradients in the loop recompute loss against a fresh graph each time, so nothing is freed between them and each backward adds to what the previous one left behind.
You run three training batches without ever calling optimizer.zero_grad(). What is the gradient applied on the third optimizer.step(), in terms of the three per-batch gradients g₁, g₂, g₃?
You built the Module abstraction in Lesson 10. PyTorch’s version adds four things: automatic parameter registration, recursive module discovery, device management and serialization. The syntax is barely different — the machinery behind it is the difference.
nn.Module is the base class for every neural network component. Subclass it, assign layers in __init__, and describe the computation in forward. That is the whole pattern, and it is the same pattern every model in the PyTorch ecosystem uses — from the three-layer network below to a 70-billion parameter transformer.
import torch
import torch.nn as nn
class MLP(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
self.layer1 = nn.Linear(input_dim, hidden_dim)
self.relu = nn.ReLU()
self.layer2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
x = self.layer1(x)
x = self.relu(x)
x = self.layer2(x)
return x
model = MLP(784, 256, 10)
pred = model(x) # __call__ → forward + hooks
params = list(model.parameters()) # recursively collected
print(sum(p.numel() for p in params)) # 200,960 + 2,570 = 203,530
# the source's 3-layer MNIST MLP (784 → 256 → 128 → 10) totals 235,146The registration is the magic worth understanding: when you assign an nn.Module or an nn.Parameter as an attribute inside __init__, PyTorch’s __setattr__ notices and registers it in the module’s parameter tree. You never call register_parameter by hand, and model.parameters() recursively collects everything — which is why one call to optimizer = Adam(model.parameters()) is enough, and why model.to("cuda") can move every tensor at once. In the mini framework you had to gather weights into a list yourself.
One convention before the arithmetic: call model(x), not model.forward(x). __call__ runs forward and the machinery around it — hooks, autocast state, mode dispatch. Direct forward() calls skip all of it.
The comparison that keeps the scale honest: GPT-2 small has about 124,000,000 parameters — 124M / 235,146 ≈ 527× the MNIST MLP. The source’s model is tiny by modern standards, and it still trains in seconds, which is exactly what makes it the right first target.
What forward actually computes. For one sample, nn.Linear(4, 3) is a matrix-vector product plus a bias. Here is the manual pass, with numbers small enough to check on paper — the same input the module lab renders:
Every layer in PyTorch is that transparent: an operation you could implement yourself, wrapped in machinery that makes it fast, movable and differentiable. On a real batch the same call runs 64 × 784 × 256 ≈ 12.8M multiply-adds for layer 1 alone — about 15.0M for the whole 3-layer stack — which is why it belongs in a kernel rather than a Python loop.
| Module | What it does | Parameters |
|---|---|---|
nn.Linear(in, out) | fully connected: y = x·Wᵀ + b | in·out + out |
nn.Conv2d(in_ch, out_ch, k) | 2-D convolution | in_ch·out_ch·k·k + out_ch |
nn.BatchNorm1d(features) | normalize activations | 2 · features |
nn.Dropout(p) | randomly zero activations | 0 |
nn.ReLU() | max(0, x) | 0 |
nn.GELU() | smooth ReLU variant | 0 |
nn.Embedding(vocab, dim) | lookup table | vocab · dim |
nn.LayerNorm(dim) | per-sample normalization | 2 · dim |
Assemble the source’s MNIST model, resize its hidden layers, and watch what model.parameters() would collect — every number is in×out + out.
nn.Linear(784, 256)(784) → (256)200,960 paramsnn.ReLU()same0 paramsnn.Dropout(0.2)same0 paramsnn.Linear(256, 128)(256) → (128)32,896 paramsnn.ReLU()same0 paramsnn.Dropout(0.2)same0 paramsnn.Linear(128, 10)(128) → (10)1,290 paramsEvery parameter is a tensor that requires grad. The module tree is also why model.to("cuda") can move everything at once — it recurses over the same tree.
A loss scores the prediction, an optimizer turns the gradient into a step, and a DataLoader feeds the batches. Then there is the loop itself — short enough to memorise, precise enough that the order of its lines is the difference between training and a silent bug.
PyTorch ships production-ready versions of everything you built by hand. The loss functions live in torch.nn:
| Loss | Task | Input it expects |
|---|---|---|
nn.MSELoss() | regression | any shape |
nn.CrossEntropyLoss() | multi-class classification | logits — not softmax |
nn.BCEWithLogitsLoss() | binary classification | logits — not sigmoid |
nn.L1Loss() | regression, robust to outliers | any shape |
nn.CTCLoss() | sequence alignment (speech, OCR) | log probabilities |
Notice the two “logits — not …” rows. They are not a footnote; they are the most common silent bug in beginner training code, and the reason CrossEntropyLoss deserves its own warning.
The optimizers live in torch.optim, and their interfaces are what you already know from Lesson 06 — you are picking defaults, not new algorithms:
| Optimizer | When to use | Typical learning rate |
|---|---|---|
SGD(params, lr, momentum) | CNNs, well-tuned pipelines | 0.01 – 0.1 |
Adam(params, lr) | default starting point | 1e-3 |
AdamW(params, lr, weight_decay) | transformers, fine-tuning | 1e-4 – 1e-3 |
LBFGS(params) | small-scale, second-order | 1.0 |
Now the loop itself. Every PyTorch training loop — every one — is this shape:
for epoch in range(num_epochs):
model.train()
for inputs, targets in train_loader:
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad() # 1. clear yesterday's gradients
outputs = model(inputs) # 2. forward: build the tape
loss = criterion(outputs, targets) # 3. score the prediction
loss.backward() # 4. reverse the tape → .grad
optimizer.step() # 5. nudge every parameter
# evaluation — same loop, tape off
model.eval()
with torch.no_grad():
for inputs, targets in test_loader:
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
...One pass, by hand. Follow the source’s autograd-demo values through all five lines — w = [1, −2, 0.5], loss = Σw², lr = 0.1:
The data pipeline. A PyTorch Dataset is an abstract class with just two methods — __len__ and __getitem__ — and DataLoader wraps it with batching, shuffling and multi-process loading. The source builds both for MNIST:
from torch.utils.data import Dataset, DataLoader
class MNISTDataset(Dataset):
def __init__(self, images, labels):
self.images = images
self.labels = labels
def __len__(self):
return len(self.labels) # 60,000
def __getitem__(self, idx):
return self.images[idx], self.labels[idx]
loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=4)The numbers that pipeline produces: ⌈60,000 / 64⌉ = 938 training batches per epoch (the last one holds 32 samples), so 10 epochs is 9,380 optimizer steps; testing uses ⌈10,000 / 256⌉ = 40 batches (last one 16). num_workers=4 spawns four background processes that prepare the next batches while the accelerator trains on the current one — on disk-bound work like large images or audio, the source notes this alone can double training speed.
Move the logits and watch softmax, the loss and its gradient p − y together. Then commit the classic mistake on purpose with the “feed probabilities” probe.
The source’s warning: “Pass raw logits, not softmax outputs. This is a common mistake that produces wrong gradients silently.” The probe shows why “silently” is the scary word — training runs, it just optimizes the wrong thing.
Run the source’s autograd demo as a real training loop. Step the five lines by hand, switch on the classic bug, and watch the loss curve tell on you.
The order is not a suggestion: zero before forward, backward before step. Five lines, in this order, trained GPT-4, Stable Diffusion and LLaMA — only the model and the data changed.
Your classification model outputs softmax probabilities, and you pass them to nn.CrossEntropyLoss(). Training runs without an error. What is wrong?
Three things stand between a correct training loop and a working one: where the tensors live, how many bits they are stored in, and whether the run can survive the end of the process. Then there is the skill nobody teaches in the forward pass — debugging.
Devices. The pattern never changes: choose the device once, move the model once, move every batch inside the loop. On a Mac, mps replaces cuda; on a laptop without a GPU, everything runs on cpu and is still correct, just slower.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
for inputs, targets in train_loader:
inputs, targets = inputs.to(device), targets.to(device)
# ...Mixed precision. Modern GPUs run float16 arithmetic much faster than float32, and storing activations in half the space doubles the batch sizes that fit in memory. PyTorch’s autocast runs the forward and backward passes in float16 while keeping float32 master weights, and GradScaler protects against the one hazard of float16: its smallest normal number is about 6.1×10⁻⁵, so tiny gradients underflow to zero. The scaler multiplies the loss by a large factor before backward (so every gradient is representable), then divides the gradients back out before the step, adjusting the factor automatically as overflow appears.
from torch.amp import autocast, GradScaler
scaler = GradScaler()
for inputs, targets in loader:
inputs, targets = inputs.to(device), targets.to(device)
with autocast(device_type="cuda"):
outputs = model(inputs)
loss = criterion(outputs, targets)
scaler.scale(loss).backward() # scale up so gradients don't underflow
scaler.step(optimizer) # unscale, then update
scaler.update() # adjust the scale factor
optimizer.zero_grad()The memory arithmetic for the 235,146-parameter MNIST model: float32 weights are 918.5 KiB, float16 weights are 459.3 KiB, and Adam’s two moments add 1.79 MiB in float32 (they stay full precision as master copies). On an A100, H100 or RTX 4090, the source expects roughly 2× throughput from autocast + GradScaler.
Checkpoints. torch.save(model.state_dict(), …) writes an OrderedDict of parameter tensors — the portable snapshot. torch.save(model, …) pickles the whole object, including a reference to your class definition, so it breaks the day you refactor the class. Load into a fresh model, then call eval():
torch.save(model.state_dict(), "model.pt")
model = MNISTModel()
model.load_state_dict(torch.load("model.pt", weights_only=True))
model.eval()Schedulers. The learning rate rarely wants to be constant. PyTorch ships 15+ schedulers — StepLR, ExponentialLR, CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau — all plugged into the same optimizer interface, so adding one costs two lines:
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=10
)
for epoch in range(10):
train_one_epoch(model, train_loader, criterion, optimizer, device)
scheduler.step() # after optimizer.step(), never beforeDebugging. When a run misbehaves, the first question is not “is my model smart enough?” but “can it learn at all?”. The smoke test is called overfit one batch: take a single batch, remove all regularization, and train on it for a few hundred steps. If the loss cannot reach ~zero, the bug is in the code — a miswired layer, a mismatched loss, a missing zero_grad(), a learning rate that is too large or too small. The debugger lab runs that experiment on 8 points that lie exactly on y = 3x − 1: a healthy run reaches w = 3, b = −1, loss ≈ 0; forgetting zero_grad() looks brilliant for two steps (the loss crashes from 17.875 to 0.198) and then rings between 0.066 and 19.7; an lr of 1.5 explodes.
PyTorch will not silently mix devices, and it will not save the model for you. Reproduce the #1 beginner error, fix it the standard way, then run the save/load round trip.
Device availability here is simulated — your browser is definitely a CPU. The error message, the fix and the state_dict rules are the real ones, verbatim from PyTorch.
The first smoke test every PyTorch engineer runs: can the model drive the loss to zero on eight points that lie exactly on a line? If not, the bug is in the code, not the data.
8 points, y = 3x − 1, model ŷ = wx + b, MSE loss, full-batch gradient descent, 200 steps. Starting at w = b = 0 the loss is 17.875; the true minimum is w = 3, b = −1 with loss 0. A simplified teaching model — the same loop in PyTorch takes a real batch and a real model.
Run the same test on your own model: one batch, a few hundred steps, no regularization. If your loss cannot reach ~zero here, more data will not help.
Likely cause. The model was moved with .to(device) but a batch was not — or vice versa.
First fix to try. device = torch.device(...); model = model.to(device); then inputs, targets = inputs.to(device), targets.to(device) inside the loop.
Likely cause. Learning rate too high, gradients exploding, or a log(0)/divide in a custom loss.
First fix to try. Lower the lr 10×, clip gradients (torch.nn.utils.clip_grad_norm_), and check for log of a zero probability — CrossEntropyLoss avoids this internally.
Likely cause. Learning rate too small, or the wrong loss for the task (MSE on logits instead of CrossEntropyLoss).
First fix to try. Raise the lr and run the source's learning-rate finder: train one epoch with exponentially increasing lr and pick the value just before the loss climbs.
Likely cause. Softmax fed into CrossEntropyLoss, labels shuffled out of sync with images, or optimizer.step() never called.
First fix to try. Pass raw logits, verify one batch by hand, and print w.grad after backward() — if it is None or zero, autograd never reached the parameters.
Likely cause. Overfitting, or model.eval() was never called so dropout is still active during evaluation.
First fix to try. Call model.eval() and wrap evaluation in torch.no_grad(); add regularization (dropout, weight decay, early stopping) if the gap persists.
You refactored your model class and now torch.load('model.pt') fails with a pickle error naming classes. What was saved, and what should have been saved instead?
Everything you built by hand has a PyTorch word for it, and the words line up one-to-one. The last step is shipping a full training run — and seeing what the same loop looks like when the framework choices multiply.
What transfers. Your Module class becomes nn.Module. Your forward method keeps its name. Your manual weight list becomes model.parameters(). Your per-layer backward() methods become one loss.backward(). Your update loop becomes optimizer.step(). Your hand-written zeroing becomes optimizer.zero_grad(). Nothing about the logic changes — the framework now supplies the autodiff, the kernels, the memory management and the serialization that you used to write yourself.
What is new. Tensors carry a device; parameters register themselves; a tape records gradients; checkpoints separate architecture from weights. Those are the four ideas this lesson added on top of Lesson 10 — and the next lesson, JAX, will show that even these are choices, not laws of nature:
| Feature | Mini framework (L10) | PyTorch | JAX |
|---|---|---|---|
| Autodiff | manual backward() | tape-based autograd | functional transforms |
| Execution | eager (Python loops) | eager (C++ kernels) | traced + JIT compiled |
| GPU support | none | CUDA · ROCm · MPS | CUDA · TPU |
| Speed (MNIST MLP) | ~300 s/epoch | ~0.5 s/epoch | ~0.3 s/epoch |
| Module system | custom Module class | nn.Module | stateless functions (Flax) |
| Debugging | print() | print(), pdb, breakpoint() | harder — JIT breaks print |
| Ecosystem | none | Hugging Face, Lightning, timm | Flax, Optax, Orbax |
| Production use | toy problems | Meta, OpenAI, Anthropic, HF | DeepMind, Midjourney |
Ship it: the full MNIST run. The source trains on MNIST without torchvision.datasets — it downloads the four gzipped IDX files, parses the binary headers with struct and gzip, flattens each 28×28 image to 784 values, and divides by 255.0 to land in [0, 1]. The training and evaluation functions are the canonical loop with bookkeeping attached:
def train_one_epoch(model, loader, criterion, optimizer, device):
model.train()
total_loss = 0
correct = 0
total = 0
for images, labels in loader:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item() * images.size(0)
_, predicted = outputs.max(1)
correct += predicted.eq(labels).sum().item()
total += labels.size(0)
return total_loss / total, correct / total
def evaluate(model, loader, criterion, device):
model.eval()
total_loss = 0
correct = 0
total = 0
with torch.no_grad():
for images, labels in loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
total_loss += loss.item() * images.size(0)
_, predicted = outputs.max(1)
correct += predicted.eq(labels).sum().item()
total += labels.size(0)
return total_loss / total, correct / total
model = MNISTModel().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)Ten epochs of that, and the source’s expected output reads:
The ecosystem is the payoff. The reason this interface matters beyond one tutorial is that every modern model library speaks it. Hugging Face’s transformers hands you a subclass of nn.Module; timm hands you vision models; PyTorch Lightning wraps this same loop in a training engine for multi-GPU runs; torch.compile can trace forward and fuse the kernels without you changing the code. The source’s production column is not decoration: Meta, OpenAI, Anthropic and Hugging Face all train on PyTorch, and when you run fine-tuning on a pretrained model, you will write the same five lines with a smaller learning rate.
The takeaway that outlasts the framework. You learned the concepts the hard way — implementing every gradient, every optimizer, every loop — so that frameworks are now tools instead of magic. When a PyTorch run misbehaves, you can reason about .grad values because you once computed them by hand. When a new framework appears (and one will), the five lines and the mental model come with you. That is what the last eleven lessons bought: you know why, so the how is just syntax.
Answer before you look. The autograd question and the CrossEntropy question are the two that separate “I followed the tutorial” from “I can debug this at 2 a.m.”.
0 / 5 answered · 0 correct
01What is autograd in PyTorch?
02What does torch.no_grad() do, and when should you use it?
03What does nn.Linear(784, 256) create in PyTorch?
04Which PyTorch method computes gradients for all parameters in the computational graph?
05Why is PyTorch's training loop significantly faster than the pure-Python mini framework?
Click a card to swap the lazy description for what it actually means.
Four ports with exact numbers — the batch-norm ledger, a custom Fashion-MNIST dataset with loader arithmetic, SGD + momentum against Adam, and the mixed-precision loop with its memory math. Try first; a worked answer is one click away.
The dropout model has three Linear layers: 784×256 + 256 = 200,960, 256×128 + 128 = 32,896, 128×10 + 10 = 1,290 — total 235,146. Each BatchNorm1d(features) adds a learnable scale γ and shift β, 2·features parameters: BatchNorm1d(256) = 512 and BatchNorm1d(128) = 256, so the batch-norm model totals 235,146 + 768 = 235,914. The buffers running_mean and running_var also exist but are not parameters — they live in state_dict, not parameters(). In training mode the layer normalizes with the current batch's mean and variance; in eval mode it uses the running estimates, which is exactly why model.eval() matters. The source's experiment (Adam, lr = 1e-3, 10 epochs, no dropout) expects 98%+ test accuracy in fewer epochs than the dropout-only run: normalization keeps each layer's input distribution stable and lets the optimizer take effective steps. Your absolute numbers will differ by seed and hardware; the parameter arithmetic above is exact.
__len__ returns len(self.labels) = 60,000; __getitem__ returns (self.images[idx], self.labels[idx]). With batch_size = 64, one epoch is ⌈60,000/64⌉ = 938 batches (937 full batches of 64 and a last batch of 32), so 10 epochs = 9,380 optimizer steps; the test set at batch_size 256 is ⌈10,000/256⌉ = 40 batches (last one 16). batch_size fixes how many samples each gradient estimate averages; shuffle=True reorders the training set every epoch so batches are not class-sorted (turn it off for evaluation to keep results comparable); num_workers=4 spawns four processes that prepare batches while the accelerator trains on the current one; drop_last=False keeps the short last batch. Fashion-MNIST has identical 28×28 greyscale format — reuse loaders, just swap the files — but clothing classes overlap visually (shirt vs pullover vs coat), so the same MLP lands around 88% instead of 98%. That gap is signal, not a bug: the task is genuinely harder.
SGD with momentum: v₁ = 0.9×0 + 0.2 = 0.2, so w₁ = 0.5 − 0.01×0.2 = 0.498; v₂ = 0.9×0.2 + 0.2 = 0.38, so w₂ = 0.498 − 0.01×0.38 = 0.4942. The step grows while the gradient stays consistent — that is the point of momentum — and eventually settles at lr·g/(1−β) = 0.01×0.2/0.1 = 0.02 per step. Adam on the same constant gradient has m̂ = g and √v̂ = |g| after bias correction, so its update is lr·1/(1 + ε/|g|) ≈ lr = 0.01 from the very first step. Compare the two philosophies: at the same lr, Adam's step is 0.01 regardless of whether g is 0.2 or 200, while SGD's settles at 0.02 here and would settle at 20 if g were 200. Adam divides the gradient's magnitude out; SGD rides it. The CosineAnnealingLR schedule anneals lr from its initial value to ~0 over T_max = 10 epochs on a half-cosine, so the last epochs take tiny steps that let SGD settle. Whether SGD catches Adam by epoch 10 depends on the task and tuning; the source's exercise is to measure it. If you try it, change the learning rate first — SGD's working range (0.01–0.1) is roughly 10–100× Adam's (1e-3).
Loop changes: wrap forward + loss in with autocast(device_type='cuda'):, then replace loss.backward() with scaler.scale(loss).backward(), optimizer.step() with scaler.step(optimizer), and add scaler.update() after each step. Memory: the weights are 235,146 × 4 = 940,584 bytes ≈ 918.5 KiB in float32 and 235,146 × 2 = 470,292 bytes ≈ 459.3 KiB in float16 — half. Adam's m and v moments are another 2 × 940,584 = 1,881,168 bytes ≈ 1.79 MiB in float32, and those stay fp32 as master copies. The scaler exists because float16's smallest normal number is about 6.1×10⁻⁵: gradients smaller than that round to zero. GradScaler multiplies the loss by a large factor S before backward, so every gradient is S times bigger and safely representable; before the optimizer step the scaler divides the gradients back out (unscales), updating S up or down automatically as overflow appears. The source expects roughly 2× throughput with autocast + GradScaler on an A100, while keeping float32 master weights for stability. If a step overflows (inf/NaN), the scaler skips it and reduces S — that is normal, not a crash.
You do not need to master these here. Each one gets a proper treatment in its own lesson; the one-line meaning is enough to keep reading. Orange dotted underlines in the prose point back to this list.
This lesson is a port of an open course. Everything here traces back to it — and the next step is running the code yourself.
Lesson text adapted from AI Engineering from Scratch (Phase 03, Lesson 11) and the Math Foundations Notebook reference build. The eight labs (dispatch race, tensor playground, autograd tape inspector, module builder, cross-entropy logits lab, five-line loop stepper, device and checkpoint simulator, and the overfit-one-batch debugger), the dispatch arithmetic behind the 500× claim (420,000 Python module calls vs 6,566 batched dispatches per epoch), the tensor memory ledger (918.5 KiB float32 → 459.3 KiB float16), the strides check that explains why .view() refuses after a permute, both fully worked autograd examples with the finite-difference check, the manual forward pass through Linear(4, 3) = 15 parameters, the cross-entropy hand check with its p − y gradient and the double-softmax probe, the 235,146-parameter ledger with its 235,914 batch-norm variant, the overfit-one-batch numbers (17.875 → ~0), and the 112 GB optimizer-state arithmetic for a 7B model are original to this page. Every number shown is computed live by the labs or verified by hand in the prose.