PyTorch
Tensors, autograd, nn.Module, and the training loop — the framework the labs actually run on. Built from zero: what a tensor really is, how PyTorch secretly records your arithmetic so it can differentiate it, and the five lines that train every model you will ever write.
Start here — what PyTorch actually is
You already know the idea of a neural network: a pile of numbers, a loss that says how wrong they are, and gradients that tell you which way to nudge them. PyTorch is the tool that does all of that bookkeeping for you.
Strip away the marketing and it does exactly three jobs:
- Hold numbers in fast multi-dimensional arrays (tensors).
- Remember every arithmetic operation you performed on them, so it can compute gradients automatically (autograd).
- Move the whole thing onto a GPU with one line, so it runs hundreds of times faster.
PyTorch is NumPy that remembers what you did to it — and can run on a GPU. That's the whole framework. Everything else is convenience wrappers around those two upgrades.
Cooking in a normal kitchen (NumPy), you chop, mix, and bake — and once it's done, nobody can tell you which step made the cake too sweet. PyTorch films every step. When the cake comes out wrong, it rewinds the tape and tells you exactly how much each ingredient contributed to the wrongness. That recording is , and it is the single reason the framework exists.
It starts from zero and builds up; flip the Depth switch at the top for the formal notation and the internals. One practical note about the code: PyTorch itself cannot run inside a browser tab, so the Run cells on this page use NumPy to rebuild PyTorch's machinery by hand — which is a better way to learn it anyway. The cells marked Colab are real PyTorch you can open and run on a free GPU.
Tensors — the container everything lives in
A is just a box of numbers with a known arrangement. That's genuinely all it is — the intimidating name is borrowed from physics, but in PyTorch it means "n-dimensional array."
You already meet them everywhere:
- A single loss value → 0 dimensions (a scalar).
- The 768 numbers describing one word → 1 dimension (a vector).
- A grayscale image, 28 rows by 28 columns → 2 dimensions.
- A colour photo: 3 colour channels × 224 rows × 224 columns → 3 dimensions.
- A batch of 32 such photos → 4 dimensions.
A 1-D tensor is a row of eggs. A 2-D tensor is an egg carton (rows of rows). A 3-D tensor is a crate of cartons. A 4-D tensor is a truck full of crates. Nothing new happens at each level — you just wrap the previous thing in another layer. The is simply the label on the truck saying how many crates, cartons, and eggs are inside.
An image really is nothing but a grid of numbers. Hover any square below to read the actual value stored in that slot — this is precisely what a model receives when you show it a picture.
The three things every tensor carries
Beyond its numbers, a tensor carries three labels. Almost every beginner bug in PyTorch is one of these three being wrong.
| Label | What it means | Typical value |
|---|---|---|
shape | How many numbers, arranged how | (32, 3, 224, 224) |
dtype | What kind of number each slot holds | torch.float32 |
device | Which piece of hardware holds it | cpu or cuda:0 |
matters more than it looks. Deep learning defaults to float32: decimals stored in 4 bytes each. Halving that to float16 halves your memory and roughly doubles your speed on modern GPUs — which is why big models are trained in half precision.
Take a batch of 32 colour photos at 224×224, stored as float32.
- Count the numbers: values.
- Each
float32takes 4 bytes. - Total: million bytes ≈ 19 MB.
Now switch the dtype to float16 (2 bytes each) and the same batch costs 9.6 MB — half. That single change is often what makes a model fit on your GPU at all.
A tensor has shape (32, 3, 224, 224). What does each of the four numbers most likely mean?
Hint: Think about a batch of colour photos.
Shapes that don't match — broadcasting
You constantly want to add a small thing to a big thing: one bias number per feature, added to every row of a batch. Writing a loop for that would be slow and ugly. does it automatically.
If you add a row of 5 numbers to a table of 4 rows by 5 columns, PyTorch mentally photocopies that row four times — once per table row — and adds. It never really makes the copies (that would waste memory); it just reads the same row four times. You get the convenience of the loop at the speed of a single bulk operation.
Run this. It builds tensors by hand and shows shapes, dtypes, and broadcasting doing its job:
You add a tensor of shape (8, 1) to a tensor of shape (1, 6). What is the resulting shape?
Autograd — the tape that remembers
Here is the part that makes PyTorch a deep-learning framework rather than a fast array library.
When you multiply and add tensors, PyTorch quietly writes down a record of what you did — a graph in which every result remembers the operation and the inputs that produced it. Call .backward() on the final number, and PyTorch walks that record in reverse, applying the chain rule at every step, and deposits a gradient on every tensor you asked it to track.
Going forwards, you compute a loss. Going backwards, PyTorch assigns blame: for each of the millions of numbers involved, how much would the loss change if I nudged this one number up a little? That blame value is the gradient, and it is exactly what the optimizer needs to know which way to turn each knob.
Imagine every operation printing a little receipt: "I am a multiply; my inputs were these two values." Stack the receipts as you cook. When the dish is wrong by some amount, you pick up the stack from the top and work down, each receipt telling you how to split the blame between its two inputs. By the bottom of the pile, every original ingredient has its share. That stack is the , and walking it downwards is .
Take and evaluate it at . Let's do what autograd does, step by step.
Forward — compute, recording each step:
Backward — start with "how much does affect itself?" The answer is 1, then push blame down:
- , so both and each get blame (nudging either by a little nudges by the same amount).
- , so gets blame .
- , so gets blame from each of the two slots it fills: .
- , so additionally gets blame .
Total blame on : . Check by hand: , so . ✓
Notice step 3 and 4 both added to 's blame rather than replacing it. That accumulation is not a quirk of this example — it is how autograd handles any value used more than once, and it has a consequence you will meet in a moment.
Now build that machine yourself. This is a real, working miniature autograd — the same idea as PyTorch's, about thirty lines:
That += in _backward is the accumulation from the worked example. Everything below is the same machine with better engineering.
The gradient autograd computes is not an approximation — it is the exact derivative. Below, the orange curve is and the teal curve is the slope autograd reports at each point. Where is falling, the slope is negative; where bottoms out, the slope crosses zero; where climbs steeply, the slope is large.
The four words you need
— the switch that says "track me."
— the receipt from the analogy.
— the rewind button.
— the answer, delivered.
Here is all four in real PyTorch:
import torch
x = torch.tensor(2.0, requires_grad=True) # 1. track this one
f = 3 * x**2 + 2 * x # 2. forward — the graph is built here
print(f.grad_fn) # <AddBackward0> — the receipt
f.backward() # 3. rewind
print(x.grad) # 4. tensor(14.) — matches our hand calculationloss.backward() adds into .grad; it does not overwrite it. Forget to clear the gradients between steps and step 2 trains on the sum of steps 1 and 2, step 3 on the sum of 1, 2 and 3, and your training quietly diverges. That is why every training loop begins with optimizer.zero_grad().
Decoding the statement p.grad += new_gradient: the += means "take what is already there and add to it," not "replace it." The behaviour is deliberate — it is what lets you split one large batch across several forward passes (gradient accumulation) — but it means clearing is your job.
You train a model and the loss decreases for a few steps, then explodes to NaN. You check and there is no zero_grad call in your loop. Why does that produce exactly this symptom?
Hint: What is in .grad on step 10 if it was never cleared?
What does calling .backward() on the loss actually do?
nn.Module — a box that owns its parameters
You could write a network as loose tensors and track each one yourself. Nobody does, because by layer four you would be passing twenty tensors around by hand. is the container that keeps a layer's numbers together with the computation that uses them.
A brick has a fixed shape and its own studs; you snap bricks together to make bigger bricks, and the assembly is itself a brick you can snap into something larger. A module is the same: nn.Linear is a module, a block of three of them is a module, and the whole network is a module. Ask any of them for .parameters() and you get every learnable number inside, however deep the nesting.
Two things, both about bookkeeping. First, model.parameters() collects every learnable tensor in the whole tree, so you can hand the lot to an optimizer in one line. Second, model.to(device) and model.state_dict() reach every one of them too — so moving to a GPU or saving a checkpoint is also one line, no matter how big the model.
import torch
import torch.nn as nn
class TinyMLP(nn.Module):
def __init__(self, in_dim=4, hidden=16, out_dim=3):
super().__init__() # always call this first
self.fc1 = nn.Linear(in_dim, hidden) # registered automatically
self.fc2 = nn.Linear(hidden, out_dim)
def forward(self, x): # you define the computation...
x = torch.relu(self.fc1(x))
return self.fc2(x) # ...PyTorch handles the backward
model = TinyMLP()
print(model) # prints the whole tree
n = sum(p.numel() for p in model.parameters())
print("learnable numbers:", n) # 4*16+16 + 16*3+3 = 131
out = model(torch.randn(8, 4)) # call the model, never model.forward()
print(out.shape) # torch.Size([8, 3])Write model(x), not model.forward(x). They look equivalent and mostly behave the same, but the direct call skips PyTorch's registered hooks — which is what silently breaks half-precision autocasting, profiling tools, and quantization down the line.
You add self.scale = torch.ones(10) to a module and it never trains, no matter what you do. What went wrong?
Hint: How does a module decide what counts as a parameter?
The training loop — the five lines that train everything
Every PyTorch training script ever written is a variation on the same five steps. Learn these and you can read any repository on GitHub.
for epoch in range(epochs):
for x, y in loader:
optimizer.zero_grad() # 1. clear last step's gradients
pred = model(x) # 2. forward — build the graph
loss = criterion(pred, y) # 3. how wrong are we? (one number)
loss.backward() # 4. backward — fill every .grad
optimizer.step() # 5. nudge every parameter downhillLine by line, in plain language:
zero_grad()— wipe the slate. Gradients accumulate, so last step's blame must go before this step's is computed.model(x)— run the data through. As a side effect, PyTorch records the graph.criterion(pred, y)— compare prediction against truth, reducing everything down to a single number. It must be a single number for step 4 to have a seed.loss.backward()— walk that graph in reverse, depositing a gradient on every parameter.optimizer.step()— read each.gradand move each parameter a small distance in the downhill direction.
Notice that no single line does two jobs. backward() computes gradients but changes nothing. step() changes parameters but computes nothing. zero_grad() only clears. This separation is why PyTorch feels like plain Python — and it is also why forgetting any one of the three fails in its own distinctive way.
Step 5 is gradient descent, the algorithm from Optimization. Drop a starting point on the surface below and watch what step() is doing, one step at a time:
Brighter = higher loss. Watch how a high learning rate overshoots, and how momentum powers through the small bumps toward a minimum.
Now watch the whole loop run. This is the five steps written out longhand in NumPy — no framework, nothing hidden — fitting a straight line to noisy data:
The learning rate decides the shape of that curve entirely. These are the real loss curves from the cell above, run at three different values of lr:
And the same fit in real PyTorch — notice how little is left to write once autograd is doing the derivatives:
import torch
import torch.nn as nn
X = torch.randn(200, 1)
y = 3.0 * X - 1.0 + 0.3 * torch.randn(200, 1)
model = nn.Linear(1, 1) # w and b live in here
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
for step in range(60):
optimizer.zero_grad()
loss = criterion(model(X), y)
loss.backward()
optimizer.step()
print(model.weight.item(), model.bias.item()) # close to 3.0 and -1.0You remove optimizer.step() from a correct training loop but leave everything else. What happens?
Data — Dataset and DataLoader
Real training does not feed the whole dataset at once. Two small classes handle the plumbing.
A answers two questions: how many examples are there, and give me number 17. A wraps it and handles the rest: grouping examples into batches, reshuffling every epoch, and prefetching on background workers.
The Dataset is the library — it knows what is on the shelves and can fetch any single book. The DataLoader is the librarian who brings you a stack of 32 books at a time, in a different random order each visit, and who has already started fetching the next stack while you are reading this one. That last part matters: without it, your GPU sits idle waiting for the disk.
from torch.utils.data import Dataset, DataLoader
class SquaresDataset(Dataset):
def __init__(self, n=1000):
self.x = torch.randn(n, 1)
self.y = self.x ** 2
def __len__(self):
return len(self.x) # how many examples
def __getitem__(self, i):
return self.x[i], self.y[i] # example number i
loader = DataLoader(SquaresDataset(), batch_size=32, shuffle=True, num_workers=2)
for xb, yb in loader: # xb: (32, 1), yb: (32, 1)
...shuffle=True on training data breaks any accidental ordering (all the cats first, then all the dogs) that would otherwise bias each batch. On validation and test data, leave it off — shuffling gains nothing and makes results harder to compare between runs.
Devices — moving onto the GPU
A tensor lives on exactly one device, and two tensors can only be combined if they live on the same one. Moving is one call.
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device) # moves in place for modules
x = x.to(device) # returns a NEW tensor for plain tensorsModules move in place, plain tensors do not. model.to(device) works; x.to(device) returns a new tensor and leaves x alone, so you must reassign it. Miss that and you get Expected all tensors to be on the same device — by far the most common GPU error there is.
Do not call .item() or print(loss) inside a tight training loop. Both force the CPU to wait for the GPU to finish, destroying the overlap that makes GPU training fast. Accumulate on the GPU and read the value once per epoch.
The errors you will actually hit
| Message | What it really means | Usual fix |
|---|---|---|
mat1 and mat2 shapes cannot be multiplied | A layer's in_features does not match what you fed it | Print x.shape just before the layer |
Expected all tensors to be on the same device | Something is on CPU, something on GPU | Reassign: x = x.to(device) |
element 0 of tensors does not require grad | The graph was broken, often by .detach(), .numpy(), or a stray no_grad | Check nothing detaches before the loss |
grad can be implicitly created only for scalar outputs | You called backward() on a non-scalar | Reduce it: loss.mean() |
| Loss is exactly flat from step 1 | Nothing is updating | Missing step(), or lr=0, or params not passed to the optimizer |
| Loss goes to NaN after a few steps | Steps far too large | Lower the learning rate; check for a missing zero_grad() |
| Loss decreases but eval accuracy is terrible | Eval running in training mode | Add model.eval() and torch.no_grad() |
When a model misbehaves, print shapes. Not values — shapes. Put print(x.shape) between every layer and run one batch. A startling fraction of "the model won't learn" turns out to be a tensor that is silently the wrong shape and being broadcast into nonsense rather than raising an error.
Explain to a friend who knows Python but not deep learning what happens between loss.backward() and optimizer.step() — using the receipt or camera picture, no formulas. Then explain why zero_grad() has to be there at all. If you stall on why gradients accumulate instead of overwrite, that is the exact spot to reread.
- PyTorch is NumPy that remembers what you did to it, and that can run on a GPU.
- A tensor is an n-dimensional box of numbers carrying three labels: shape, dtype, and device. Most beginner bugs are one of those three being wrong.
- Autograd records every operation into a graph as you compute forwards, then walks it backwards applying the chain rule, depositing a gradient in every parameter's
.grad. - Gradients accumulate rather than overwrite — which is why
zero_grad()is not optional. - nn.Module bundles parameters with computation, so
.parameters(),.to(device)and.state_dict()reach everything however deeply nested. - The training loop is always the same five lines:
zero_grad→ forward → loss →backward→step. Each does exactly one job.
Practice — and how to make it stick
• Retrieval practice: before scrolling back, try to write the five-line training loop from memory. Struggling to recall it beats rereading it ten times.
• Spaced repetition: mark this topic complete to add it to your Review queue, so it resurfaces right before you would have forgotten it.
• Interleaving: mix these exercises with Backpropagation and Optimization problems rather than doing them in a block — messier practice, sturdier memory.
- Break it on purpose. Take the working loop above and remove
zero_grad(), thenstep(), then flip the minus sign to a plus. Predict what each will do before you run it, then check. Deliberately causing the three classic failures is the fastest way to recognize them later. - Extend the mini-autograd. Add
__pow__and atanhmethod to theValueclass, then use it to compute the gradient of a two-input function and check it against the derivative you work out by hand. - From scratch, then with the framework. Implement logistic regression twice — once in raw NumPy with hand-derived gradients, once with
nn.Moduleand an optimizer. Confirm they reach the same weights. - Read a real repository. Open any training script on GitHub and find the five lines. They are always there, sometimes with a scheduler or gradient clipping wedged between
backward()andstep().
Try it right here — edit and run the code, and if you get stuck or hit an error, ask Ada on the right: she can see your code and terminal output.
Next: see exactly what autograd is doing under the hood in Backpropagation, then put the loop to work in Regularization.