Knowledge BaseFundamentals

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.

intermediate#pytorch#autograd#training-loop#tensors

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:

  1. Hold numbers in fast multi-dimensional arrays (tensors).
  2. Remember every arithmetic operation you performed on them, so it can compute gradients automatically (autograd).
  3. Move the whole thing onto a GPU with one line, so it runs hundreds of times faster.
The one-sentence version

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.

Think of it like a kitchen with a camera in the ceiling:

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.

How to read this page — and about the code cells

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.
Think of it like nested boxes:

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.

A tensor of shape 8 by 8 — an image, seen as the numbers it actually is— interactive, drag & zoom
Loading chart…
Every cell holds one float between 0 (black) and 1 (white). To you it looks like a handwritten 7; to the model it is 64 numbers in a grid. Hover to read them. A real photo is the same thing at shape 3 by 224 by 224.

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.

LabelWhat it meansTypical value
shapeHow many numbers, arranged how(32, 3, 224, 224)
dtypeWhat kind of number each slot holdstorch.float32
deviceWhich piece of hardware holds itcpu 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.

How much memory does one batch cost?

Take a batch of 32 colour photos at 224×224, stored as float32.

  1. Count the numbers: 32×3×224×224=4,816,89632 \times 3 \times 224 \times 224 = 4{,}816{,}896 values.
  2. Each float32 takes 4 bytes.
  3. Total: 4,816,896×419.34{,}816{,}896 \times 4 \approx 19.3 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.

Try to recall

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.

Stretch the small one to fit, for free

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:

Python · runs in your browser
What this does: Builds arrays the way PyTorch builds tensors, prints their shape and dtype, and shows broadcasting adding one row of 5 biases to all 4 rows of a batch at once. Try changing bias to np.arange(4) — you will get a shape error, which is broadcasting refusing to guess.

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.

Forwards you compute, backwards it blames

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.

Think of it like a receipt for every step of the meal:

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 .

Blame assignment by hand

Take f(x)=3x2+2xf(x) = 3x^2 + 2x and evaluate it at x=2x = 2. Let's do what autograd does, step by step.

Forward — compute, recording each step:

  1. a=x×x=2×2=4a = x \times x = 2 \times 2 = 4
  2. b=3×a=3×4=12b = 3 \times a = 3 \times 4 = 12
  3. c=2×x=2×2=4c = 2 \times x = 2 \times 2 = 4
  4. f=b+c=12+4=16f = b + c = 12 + 4 = 16

Backward — start with "how much does ff affect itself?" The answer is 1, then push blame down:

  1. f=b+cf = b + c, so both bb and cc each get blame 11 (nudging either by a little nudges ff by the same amount).
  2. b=3ab = 3a, so aa gets blame 3×1=33 \times 1 = 3.
  3. a=x×xa = x \times x, so xx gets blame x×3x \times 3 from each of the two slots it fills: 2×3+2×3=122 \times 3 + 2 \times 3 = 12.
  4. c=2xc = 2x, so xx additionally gets blame 2×1=22 \times 1 = 2.

Total blame on xx: 12+2=1412 + 2 = \mathbf{14}. Check by hand: f(x)=6x+2f'(x) = 6x + 2, so f(2)=14f'(2) = 14. ✓

Notice step 3 and 4 both added to xx'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:

Python · runs in your browser
What this does: A miniature autograd in ~30 lines — the same design PyTorch uses. Each Value remembers which Values produced it and how to pass blame back to them, so calling backward() walks that record in reverse and fills in every gradient. It computes f = 3x² + 2x at x = 2 and checks its answer against the derivative done by hand (6x + 2 = 14). Try changing the expression, or x, and see it still get the right answer.

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 ff and the teal curve is the slope autograd reports at each point. Where ff is falling, the slope is negative; where ff bottoms out, the slope crosses zero; where ff climbs steeply, the slope is large.

What autograd computes — the function and its exact slope— interactive, drag & zoom
Loading chart…
Orange: f(x) = 3x² + 2x. Teal: the gradient autograd returns, f'(x) = 6x + 2. The gradient crosses zero exactly where the curve flattens out, at x = -1/3 — which is why an optimizer that follows the gradient downhill comes to rest there.

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:

Python · needs a GPU — run on Colab
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 calculation
The bug that bites literally everyone

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

Try to recall

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.

Think of it like Lego bricks that carry their own studs:

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.

What the class actually buys you

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.

Python · needs a GPU — run on Colab
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])
Call the model, not forward

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.

Try to recall

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.

Python · needs a GPU — run on Colab
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 downhill

Line by line, in plain language:

  1. zero_grad() — wipe the slate. Gradients accumulate, so last step's blame must go before this step's is computed.
  2. model(x) — run the data through. As a side effect, PyTorch records the graph.
  3. 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.
  4. loss.backward() — walk that graph in reverse, depositing a gradient on every parameter.
  5. optimizer.step() — read each .grad and move each parameter a small distance in the downhill direction.
The division of labour

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:

Gradient Descent— interactive, try itOpen in lab →
Click anywhere to drop a new starting point.

Brighter = higher loss. Watch how a high learning rate overshoots, and how momentum powers through the small bumps toward a minimum.

Every call to optimizer.step() is one of these arrows. The learning rate is the arrow length; loss.backward() is what computes its direction.

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:

Python · runs in your browser
What this does: The complete five-step training loop written by hand: forward, loss, backward, step, repeat. It fits a line to noisy data generated from w=3, b=-1 and should recover roughly those values, while the plotted loss curve falls toward zero. Try lr = 0.5 (much faster) or lr = 0.9 (watch it blow up) — the same learning-rate behaviour you get with a real optimizer.

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:

Same model, same data, three learning rates— interactive, drag & zoom
Loading chart…
Measured by running the loop above with lr = 0.5, 0.1 and 0.02. Note the log scale on the loss. A bold rate converges in three steps; a timid one is still far from the floor after thirty. Every model you train needs this dial found by experiment — it is the first thing to tune and the first suspect when training misbehaves.

And the same fit in real PyTorch — notice how little is left to write once autograd is doing the derivatives:

Python · needs a GPU — run on Colab
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.0

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

Think of it like a library and a librarian:

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.

Python · needs a GPU — run on Colab
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 the training set, not the validation set

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.

Python · needs a GPU — run on Colab
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 tensors
The two device traps

Modules 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

MessageWhat it really meansUsual fix
mat1 and mat2 shapes cannot be multipliedA layer's in_features does not match what you fed itPrint x.shape just before the layer
Expected all tensors to be on the same deviceSomething is on CPU, something on GPUReassign: x = x.to(device)
element 0 of tensors does not require gradThe graph was broken, often by .detach(), .numpy(), or a stray no_gradCheck nothing detaches before the loss
grad can be implicitly created only for scalar outputsYou called backward() on a non-scalarReduce it: loss.mean()
Loss is exactly flat from step 1Nothing is updatingMissing step(), or lr=0, or params not passed to the optimizer
Loss goes to NaN after a few stepsSteps far too largeLower the learning rate; check for a missing zero_grad()
Loss decreases but eval accuracy is terribleEval running in training modeAdd model.eval() and torch.no_grad()
The single most useful debugging habit

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 it yourself

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.

Recap — the key ideas
  • 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 → backwardstep. Each does exactly one job.

Practice — and how to make it stick

Learn it the way that actually works


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.

  1. Break it on purpose. Take the working loop above and remove zero_grad(), then step(), 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.
  2. Extend the mini-autograd. Add __pow__ and a tanh method to the Value class, then use it to compute the gradient of a two-input function and check it against the derivative you work out by hand.
  3. From scratch, then with the framework. Implement logistic regression twice — once in raw NumPy with hand-derived gradients, once with nn.Module and an optimizer. Confirm they reach the same weights.
  4. 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() and step().

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.

Practice lab
Your task: Prove to yourself that gradients accumulate. The starter calls backward() three times on the same expression without ever clearing x.grad. Run it and predict what you will see before reading the output. Then do the TODO: add a line that resets x.grad to 0.0 at the top of each loop iteration — the by-hand version of optimizer.zero_grad() — and confirm every iteration now reports the same correct gradient of 14.
editor
terminal
Press Run (⌘/Ctrl+Enter) to execute.
Ask Ada — she can read your terminal

Next: see exactly what autograd is doing under the hood in Backpropagation, then put the loop to work in Regularization.

Key papers