Lesson 4: PyTorch basics — Tensor & Autograd in depth

In Lesson 3 we got comfortable with vectorisation and broadcasting in NumPy. But NumPy only runs on the CPU, and it cannot compute derivatives for you. To train the large deep-learning architectures ahead — Transformers, deep CNNs — we need a more capable tool: PyTorch.

This lesson introduces PyTorch through its core data structure, the Tensor; shows how to control the shape of your data with reshaping; and then digs into Autograd, the automatic differentiation engine that makes backpropagation possible.

4.1 What is a Tensor, and how does it differ from a NumPy array?

Mathematically and in terms of memory layout, a PyTorch Tensor is 99% the same thing as a NumPy ndarray. It is a multi-dimensional array holding elements of one type, laid out contiguously in RAM. What PyTorch adds are two capabilities built specifically for deep learning:

1. Hardware acceleration: a PyTorch tensor can move its computation from the CPU to a high-performance GPU — NVIDIA cards through CUDA, or Apple Silicon through MPS — so that millions of matrix operations run in parallel.
2. Automatic differentiation (Autograd): every tensor can record its own computation history and produce the gradient at any point in that history on demand.

Before computing anything, you have to tell PyTorch where the tensor lives: in ordinary RAM (CPU) or in the graphics card's memory. The snippet below does two things — detect what acceleration your machine has, then move the tensor there using the device attribute:

import torch

# A tensor starts life on the CPU unless you say otherwise.
x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

# Pick the best hardware available, in order of preference.
device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using device: {device}")

# .to() returns a COPY on the target device. It does not move x in place.
x_gpu = x.to(device)
print(x_gpu)

That chained if...else is Python's equivalent of nested ternaries in JavaScript, and the order matters: cuda is a discrete NVIDIA card (fastest), mps is Apple Silicon on a Mac, and cpu is the fallback that always works. On an M-series MacBook you will see Using device: mps; on a machine with no GPU you get cpu — and the lesson still runs fine, just slower.

One detail is easy to misread, on the last line: x.to(device) does not move x. It returns a new copy on the target device, and the original x stays on the CPU. Forgetting to assign the result — writing x.to(device) and then carrying on with x — is the fastest route into the pitfall immediately below.

Because both store their data contiguously in the same way, PyTorch lets you convert between a CPU tensor and a NumPy array almost for free, via .numpy() and torch.from_numpy() — but "almost for free" hides an important trap:

import torch
import numpy as np

t = torch.ones(3)
n = t.numpy()  # No copy — n and t point at the same block of RAM.

n[0] = 99.0  # Change the NumPy array...
print(t)     # ...and the tensor changed too: tensor([99., 1., 1.])

# Want them independent? Copy on purpose.
n_independent = t.numpy().copy()
n_independent[0] = -1.0
print(t)  # Unchanged this time: tensor([99., 1., 1.])

If you come from JavaScript, this behaves exactly like assigning an object: const b = a creates no new object, so changing b.x also changes a.x. What makes it dangerous here is that the two variables belong to two different libraries and look entirely unrelated — so when the data goes wrong at the end of a pipeline, almost nobody suspects this.

4.2 Reshaping tensors (view vs reshape)

We have seen what a tensor holds and where it lives. This section is about changing its shape — something you will do at almost every layer of a neural network from Lesson 5 onwards.

"Shape" is the size of each dimension: a 28×28 greyscale image has shape (28, 28), and a batch of 64 such images has shape (64, 28, 28). Every layer expects its input in a particular shape, so folding data to fit happens constantly. PyTorch offers two methods for it, .view() and .reshape() — they look interchangeable, but differ in one way that trips up newcomers with a runtime error.

The theory only really lands when you watch it crash. The code below does one thing: reshape the same tensor to (6, 1) twice — once while it is still contiguous, once after a transpose has made it non-contiguous:

import torch

a = torch.tensor([[1, 2, 3], [4, 5, 6]])  # shape (2, 3), contiguous
print(a.is_contiguous())  # True

b = a.view(6, 1)  # Fine: no copy, just a different way of indexing the same RAM.

a_t = a.t()  # Transpose -> shape (3, 2), but the RAM layout did NOT change.
print(a_t.is_contiguous())  # False

try:
    a_t.view(6, 1)
except RuntimeError as e:
    print("Error:", e)
    # Error: view size is not compatible with input tensor's size and stride...

c = a_t.reshape(6, 1)               # Option 1: reshape copies when it has to.
d = a_t.contiguous().view(6, 1)     # Option 2: make it contiguous yourself first.

The two print lines are the whole lesson: same data, yet a.is_contiguous() returns True while a_t.is_contiguous() returns False. The transpose never touched RAM — it only changed how PyTorch interprets that RAM — and since .view() needs the data laid out contiguously, it refuses to run.

Two more helpers let you add and remove size-1 dimensions:

• .squeeze(): removes every dimension of size 1. A tensor of shape (1, 5, 1) becomes (5,).
• .unsqueeze(dim): inserts a size-1 dimension at position dim. A tensor of shape (5,) with .unsqueeze(1) becomes (5, 1).

These sound like trivia, but you will use them constantly, because nearly every PyTorch layer is written to process a batch of samples at once rather than a single sample. When you want a prediction for exactly one image of shape (28, 28), the model rejects it because it expects (batch, 28, 28) — and .unsqueeze(0) is how you add that "batch of one" dimension. Going the other way, .squeeze() strips the leftover wrapper off a result before you show it to a user.

4.3 Automatic differentiation (Autograd and the computation graph)

Everything we have done with tensors so far, NumPy could also do. This section is the reason PyTorch exists — and the reason you will never hand-derive gradients for your own network.

A reminder from Lesson 2: to train a model we need to know which direction to nudge each weight so the error goes down, and that answer is the partial derivative of the error with respect to that weight. In Lesson 2 we differentiated a one-variable function by hand. A real neural network has millions of weights and hundreds of nested operations, which makes hand-differentiation impossible. PyTorch solves this with the Autograd engine.

When you create a tensor with requires_grad=True, PyTorch starts recording every operation applied to it. During the forward pass it builds a dynamic computation graph — a directed acyclic graph (DAG) of those operations.

Every value produced by an operation carries a grad_fn attribute pointing back at the operation that created it (for example <PowBackward0> or <AddBackward0>). Leaf nodes — the tensors you created yourself, which are not the result of any operation — have grad_fn = None.

Consider this simple computation graph:

The diagram below is that same graph, drawn the way PyTorch builds it in memory. The top row is the forward pass that happens when you write those two lines of code; the bottom row is the backward pass that z.backward() runs for you:

When you call z.backward(), PyTorch walks the graph backwards from $z$ to $x$, applying the chain rule from calculus:

The most important part of the diagram is the grad_fn labels on the top row. You never wrote a line to create them — PyTorch attaches to each intermediate result a pointer back to the operation that produced it. That chain of pointers is the graph, and it is what lets the backward pass know which route to take. x alone has grad_fn = None, because you created it rather than computing it — which is the definition of a leaf node.

import torch

# requires_grad=True turns on the recording. Without it, x.grad stays None.
x = torch.tensor(3.0, requires_grad=True)

# Forward pass: computes the values AND builds the graph at the same time.
y = x ** 2
z = 2 * y + 3
print(y.grad_fn)  # <PowBackward0 object at 0x...>
print(x.grad_fn)  # None — x is a leaf, nothing produced it

# Backward pass: walk the graph in reverse, applying the chain rule.
z.backward()

# Expected: dz/dx = 4x = 4 * 3 = 12
print(f"dz/dx at x=3 is: {x.grad.item():.1f}")  # dz/dx at x=3 is: 12.0

The three print lines confirm exactly what the diagram describes: y carries a PowBackward0 pointer back to the squaring that produced it, x carries nothing because it is a leaf, and after .backward() the value 12 sits in x.grad. You wrote no derivative formula at all — only the forward computation, and PyTorch worked out the rest.

Another runtime error beginners hit constantly: calling .backward() on a tensor that is not a scalar:

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x * 2  # y holds 3 values, not a single number

try:
    y.backward()  # Fails
except RuntimeError as e:
    print("Error:", e)
    # Error: grad can be implicitly created only for scalar outputs

# Fix 1: reduce to a scalar first. This is what every real loss function does,
# via .sum() or .mean().
y.sum().backward()
print(x.grad)  # tensor([2., 2., 2.])

# Fix 2: if you genuinely need per-element gradients, pass a weight vector with
# the same shape as y.
x.grad.zero_()
y2 = x * 2
y2.backward(torch.tensor([1.0, 1.0, 1.0]))
print(x.grad)  # tensor([2., 2., 2.]) — same as .sum().backward() here

This is why every real loss function (MSE, cross-entropy, and the rest) returns exactly one scalar representing the whole batch. It is not an arbitrary convention: .backward() only knows how to propagate backwards from a single starting point.

4.4 Hands-on project: fitting a polynomial with Autograd

To put Autograd to work, we will build a small optimiser of our own: fitting a cubic polynomial to a sine curve.

The polynomial we want to learn is:

where $w_1, w_2, w_3, b$ are the parameters that get nudged against the gradient, over and over, to reduce the error.

Why this problem? Because it is small enough to check by eye, yet it contains every part of a real training loop: noisy data, learnable parameters, a loss function, gradients, and an update step. In Lesson 5 you will swap those four numbers for the thousands of weights in a neural network — but the loop will still be the same five steps you see below.

Before reading the code, three terms in it are worth pinning down:

• Epoch — one full pass over the training data. We run 2,000 epochs, so the dataset goes through the model 2,000 times, adjusting the weights a little each time.
• Learning rate — how far each update travels in the direction the gradient points. You met it in Lesson 2: too large and you overshoot the minimum and the error diverges to infinity; too small and you head the right way but far too slowly to converge.
• Batch — how many samples go through the model at once. Here we use all 2,000 in one go, so the batch is the entire dataset; later lessons split it up once the data no longer fits in memory.

You can download the complete Python file from the box below:

import torch
import math

# Fit a cubic polynomial to sin(x) using nothing but Autograd.
#
#   y_pred = w1*x + w2*x^2 + w3*x^3 + b
#
# There is no nn.Module and no optimizer here on purpose: every weight update is
# written by hand, so you can see exactly what PyTorch does for you later on.


def generate_data(num_samples=2000):
    # x spread evenly across [-pi, pi]
    x = torch.linspace(-math.pi, math.pi, num_samples, dtype=torch.float32)
    # The 0.1 * randn term is Gaussian noise. Real measurements are never clean,
    # and fitting perfectly clean data teaches the wrong lesson about overfitting.
    y = torch.sin(x) + 0.1 * torch.randn(num_samples)
    return x, y


def train_autograd():
    x, y = generate_data()

    # requires_grad=True is the whole trick: from now on PyTorch records every
    # operation these four tensors take part in, so it can differentiate later.
    w1 = torch.randn((), dtype=torch.float32, requires_grad=True)
    w2 = torch.randn((), dtype=torch.float32, requires_grad=True)
    w3 = torch.randn((), dtype=torch.float32, requires_grad=True)
    b = torch.randn((), dtype=torch.float32, requires_grad=True)

    # Hyperparameters. 1e-6 looks tiny, but the loss below is a SUM over 2000
    # samples, so each gradient is roughly 2000x larger than a per-sample one.
    learning_rate = 1e-6
    epochs = 2000

    print("=== Fitting a cubic to sin(x) with PyTorch Autograd ===")
    print(
        f"Initial weights: w1={w1.item():.4f}, w2={w2.item():.4f}, "
        f"w3={w3.item():.4f}, b={b.item():.4f}\n"
    )

    for epoch in range(1, epochs + 1):
        # Forward pass — this line also builds the computation graph.
        y_pred = w1 * x + w2 * (x**2) + w3 * (x**3) + b

        # Squared error, summed. backward() needs a single number to start from.
        loss = (y_pred - y).pow(2).sum()

        # Backward pass: walk the graph in reverse and fill in every .grad.
        loss.backward()

        # The update itself is plain arithmetic, not part of the model, so keep
        # it out of the graph.
        with torch.no_grad():
            w1 -= learning_rate * w1.grad
            w2 -= learning_rate * w2.grad
            w3 -= learning_rate * w3.grad
            b -= learning_rate * b.grad

            # Gradients ACCUMULATE by default. Skip this and epoch 2 optimises
            # using epoch 1 + epoch 2 added together, which points nowhere useful.
            w1.grad.zero_()
            w2.grad.zero_()
            w3.grad.zero_()
            b.grad.zero_()

        if epoch % 200 == 0:
            print(f"Epoch {epoch:4d} | Loss: {loss.item():.4f}")

    print("\n=== Result ===")
    print(
        f"Learned polynomial: y_pred = {w1.item():.4f}*x + {w2.item():.4f}*x^2 "
        f"+ {w3.item():.4f}*x^3 + {b.item():.4f}"
    )
    print("Target function:    y = sin(x)")

    # Report the MEAN squared error, not the sum, so the number is comparable
    # across different dataset sizes.
    with torch.no_grad():
        final_y_pred = w1 * x + w2 * (x**2) + w3 * (x**3) + b
        final_loss = (final_y_pred - y).pow(2).mean()
        print(f"Final mean squared error: {final_loss.item():.6f}")


if __name__ == "__main__":
    train_autograd()

Running it and reading the output

Run python autograd_estimation.py and you will see roughly this (the numbers vary a little between runs, because the weights start at random values and the data is noisy):

=== Fitting a cubic to sin(x) with PyTorch Autograd ===
Initial weights: w1=-0.6197, w2=-0.6255, w3=-0.5777, b=-1.0142

Epoch  200 | Loss: 1255.4491
Epoch  400 | Loss: 586.2173
Epoch  600 | Loss: 282.7377
Epoch  800 | Loss: 144.6042
Epoch 1000 | Loss: 81.4823
Epoch 1200 | Loss: 52.5188
Epoch 1400 | Loss: 39.1717
Epoch 1600 | Loss: 32.9937
Epoch 1800 | Loss: 30.1211
Epoch 2000 | Loss: 28.7794

=== Result ===
Learned polynomial: y_pred = 0.8384*x + 0.0048*x^2 + -0.0907*x^3 + -0.0326
Target function:    y = sin(x)
Final mean squared error: 0.014387

Three things in this output are worth reading closely, and each says something about how models learn.

First, the error falls fast and then slows down. Between epoch 200 and 400 the loss more than halves (1,255 → 586). Between epoch 1,800 and 2,000 it only drops from 30.1 to 28.8. That is the characteristic shape of gradient descent: far from the minimum the slope is steep, so each step covers a lot of ground; close to it the slope flattens and the steps get shorter. The loss is still falling at epoch 2,000 — the model has not fully converged, and more epochs would improve it a little further.

Second, $w_2$ and $b$ both land near zero while $w_1$ and $w_3$ do not. That is not chance; the model is discovering a mathematical property of the target function. $\sin$ is an odd function: $\sin(-x) = -\sin(x)$, so its graph is symmetric about the origin. The odd-power terms ($x$ and $x^3$) share that property; the even-power term ($x^2$) and the constant do not. The only way to fit an odd function is to suppress both even components — which is exactly what Autograd works out over 2,000 rounds, without anyone telling it that $\sin$ is odd. Run it repeatedly and you will see $w_2$ always sitting around $\pm 0.01$ while $w_3$ stays around $-0.09$.

Third, is a final error of 0.0144 good or bad? Compare it against the noise we added to the data ourselves: 0.1 * torch.randn(...), a standard deviation of 0.1, so a variance of $0.1^2 = 0.01$. In other words, even a perfect model that knew $\sin(x)$ exactly could not score below about 0.01 on this dataset, because the noise is unpredictable by construction. We reached 0.0144 — only slightly above that theoretical floor. The cubic has learned nearly everything there is to learn.

Summary and what comes next

Bridge to the next lesson: from these individual tensor building blocks, Lesson 5 assembles them into a complete multi-layer neural network (MLP).