PyTorch for Beginners: Tensors, Autograd and the Standard Training Loop

The short answer

PyTorch is one of the mainstream deep-learning frameworks, and most code exercises on this site (nanoGPT, TRL and others) use it. Learning it comes down to two ideas: tensors record how they were computed, and loss.backward() computes the gradient of every parameter for you — plus one five-step training loop. This lesson belongs to the PyTorch & Autograd unit; the prerequisite is Python and NumPy.

Step 1: tensors and devices

Tensors behave almost exactly like NumPy arrays: they have a shape and a dtype and support broadcasting and matrix maths. The differences: they can run on a GPU, and they support automatic differentiation (official tutorial).

import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
a = torch.randn(3, 4, device=device)
b = torch.randn(4, 2, device=device)
c = a @ b            # matrix multiply, shape (3, 2)

Tensors in the same operation must be on the same device — one of the most common beginner errors. Install with the selector on the official Get Started page for your OS and CUDA version rather than a random pip install.

Step 2: how autograd computes gradients

Training a network means repeatedly "compute gradients, update parameters". Deriving every layer by hand is slow and error-prone. Automatic differentiation hands it to the program: record each operation during the forward pass, then apply the chain rule step by step in reverse. This reverse mode is the general form of backpropagation in neural networks (Baydin et al., 2015).

In PyTorch you just mark tensors with requires_grad=True:

import torch
x = torch.tensor(2.0, requires_grad=True)
w = torch.tensor(3.0, requires_grad=True)
y = w * x ** 2 + 1          # y = 3·4 + 1 = 13
y.backward()
print(x.grad, w.grad)       # dy/dx = 2wx = 12, dy/dw = x² = 4
print(y.grad_fn)            # <AddBackward0 ...>: y came from an addition

Key points from the docs (autograd tutorial, Autograd mechanics):

• During the forward pass autograd records a directed acyclic graph (DAG); each result's grad_fn points to the operation that produced it.
• backward() starts from the result and applies the chain rule down to the leaf tensors (the parameters you created), storing gradients in their .grad.
• Graphs are dynamic: after each backward() the next forward pass rebuilds the graph from scratch, so models can use Python if and for freely. The PyTorch paper calls this style imperative — the model is a regular Python program (Paszke et al., 2019).
• Gradients accumulate by default; nothing resets them for you:

w = torch.tensor(1.0, requires_grad=True)
for i in range(2):
    (w * 3).backward()
    print(w.grad)           # 3 the first time, 6 the second: accumulated

The best way to see what autograd does under the hood is to build one yourself with Karpathy's micrograd: about 100 lines for the autograd engine and 50 for a PyTorch-like neural-net library.

Step 3: the standard training loop

Rewrite last lesson's NumPy linear regression in PyTorch and the hand-written gradient lines disappear:

import torch
from torch import nn

torch.manual_seed(0)
X = torch.rand(200, 1)
y = 3 * X + 2 + 0.1 * torch.randn(200, 1)

model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
opt = torch.optim.SGD(model.parameters(), lr=0.5)

for step in range(500):
    pred = model(X)              # 1. forward
    loss = loss_fn(pred, y)      # 2. loss
    opt.zero_grad()              # 3. clear last step's gradients
    loss.backward()              # 4. backward
    opt.step()                   # 5. update parameters

print(model.weight.item(), model.bias.item())   # close to 3 and 2

These five steps are the optimisation loop from the official tutorial, which notes that gradients add up by default, so you zero them explicitly each iteration to prevent double-counting (Optimizing Model Parameters). Pre-training and fine-tuning large models are the same five steps — with bigger models, more data and many GPUs.

Step 4: a fuller version (data loading, Adam, evaluation)

Real projects feed data in mini-batches and usually use Adam, which adapts each parameter's step size using estimates of the gradient's first and second moments (Kingma & Ba, 2014).

import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader

device = "cuda" if torch.cuda.is_available() else "cpu"
X = torch.randn(1000, 20)
y = (X[:, 0] + X[:, 1] > 0).long()          # a simple binary task
loader = DataLoader(TensorDataset(X, y), batch_size=64, shuffle=True)

model = nn.Sequential(
    nn.Linear(20, 64), nn.ReLU(), nn.Dropout(0.1), nn.Linear(64, 2)
).to(device)
loss_fn = nn.CrossEntropyLoss()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(5):
    model.train()                            # training mode: Dropout active
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        loss = loss_fn(model(xb), yb)
        opt.zero_grad()
        loss.backward()
        opt.step()

    model.eval()                             # eval mode: Dropout off
    with torch.no_grad():                    # don't record the graph
        acc = (model(X.to(device)).argmax(1) == y.to(device)).float().mean().item()
    print(epoch, round(acc, 3))

Notes:

• Dataset and DataLoader: a Dataset answers "what is sample i"; a DataLoader batches, shuffles and loads in parallel (data tutorial).
• model.train() / model.eval(): Dropout and BatchNorm behave differently in training and inference; the official tutorial switches modes in its train and test functions.
• torch.no_grad(): evaluation and inference need no gradients; inside this block no graph is recorded, saving memory and time.
• Reproducibility: set torch.manual_seed, but fully identical results are not guaranteed across releases, platforms, or CPU vs GPU (Reproducibility).

Common pitfalls

Try it yourself

1. Train your first image classifier with the official 60-minute blitz.
2. Solve the 21 Tensor Puzzles, implementing common tensor functions with broadcasting and basic ops only.
3. Unit project: build a ~100-line autograd engine from scratch following micrograd, then write the same two-layer network in PyTorch and check that both produce the same gradients.
4. For a runnable walkthrough, read the automatic differentiation section of Dive into Deep Learning.

Next, move on to How Neural Networks Work to see what each layer is doing.