PyTorch, from first principles

PyTorch tensors support shapes and broadcasting much like NumPy arrays. PyTorch can also record operations to calculate derivatives automatically. This is autograd. We’ll use it to fit a line by repeatedly calculating a loss, finding its gradient, and updating the parameters.

Data, device, and gradients

A tensor’s device identifies where its data lives, such as CPU or GPU memory. In ordinary gradient mode, differentiable operations involving an input with requires_grad=True are recorded for backward calculation. no_grad and inference mode disable this recording.

Compare the tensor’s values, device, and gradient history.

Try it

run the op with the memory switch off — the result is just a number, history discarded, like NumPy. Flip the switch on and run it again: the same op now leaves a trail of grad_fn chips behind it. The recorded operations are what backward uses to calculate gradients.

Why remember? Because slopes multiply

Recall what the math book ended on: learning is gradient descent — every parameter needs its slope on the loss, and the chain rule says slopes multiply through composed steps. A loss is always a composition: multiply, add, activate, compare. So to get every slope, you need to know which steps made the loss, in what order.

The computational graph records these dependencies. Autograd uses them to calculate derivatives:

• Forward: calculate the loss and record the operations needed for differentiation.
• loss.backward(): apply the chain rule backward, multiplying along paths and adding where they meet. Parameter gradients accumulate in .grad.

Forward writes the tape left to right. backward() replays it right to left, multiplying local slopes as it goes, and the products land in .grad.

Try it

press backward() and follow the running product — every stop is one local slope, nothing more. Then press step: the weights slide against their gradients and the loss drops from 9.0 to 2.25. You just watched backpropagation and one stride of gradient descent, with nothing hidden.

The training loop

An optimizer updates the model’s parameters using their gradients. Basic gradient descent uses w ← w − lr·grad; other optimizers use different update rules. Here is the loop for this regression example:

for epoch in range(200):
    pred = model(x)
    loss = F.mse_loss(pred, y)
    opt.zero_grad()
    loss.backward()
    opt.step()

The loop on the left, its consequences on the right: each pass nudges the line toward the data, and the loss curve is the ball from the gradient chapter, rolling.

Try it

run one epoch (slow) and watch each line light as it acts. Then crank the learning rate past 1 and train — the fit line starts vaulting the data instead of settling onto it. You’ve seen this exact failure before, as a ball overshooting a valley.

The math book, running

The loop connects several ideas from the earlier lessons:

• loss = ... — this regression example uses F.mse_loss. Classification models often use F.cross_entropy, connected to average surprise.
• training on batches — the sampling license: the batch gradient is a noisy estimate whose expectation is the true one.
• loss.backward() — the chain rule with a tape, this article.
• opt.step() with its learning rate — the stride.

Go deeper: why the tape is rebuilt every pass

Eager-mode autograd builds a fresh graph as each forward pass runs, including the branches taken by Python control flow. By default, backward frees saved intermediate values. Another backward pass through the same calculation may therefore fail unless you retain the graph or recompute the forward pass. See PyTorch’s autograd explanation.

The example uses three core operations: calculate a loss, differentiate it, and update the parameters. Larger programs also use nn.Module to organize model components and parameters, and DataLoader to provide batches of data.