PyTorch I: tensors, autograd and the training loop

Tensors: NumPy arrays that remember where they came from

A torch.Tensor is an n-dimensional array with the same shape/dtype/broadcasting semantics as NumPy (2.2 NumPy and vectorized thinking) plus two additions: it can live on an accelerator, and it can record the operations that produced it so that gradients can be computed. Conversion is cheap in both directions: torch.from_numpy(a) shares memory with the NumPy array on CPU; t.numpy() goes back (after .detach().cpu() if the tensor has gradients or lives on a GPU).

import torch
device = "cuda" if torch.cuda.is_available() else "cpu"   # "mps" on Apple silicon
x = torch.tensor([[1., 2.], [3., 4.]])            # float32 by default (not float64!)
y = torch.arange(6).reshape(2, 3)                 # int64
x = x.to(device)                                  # move data; the model must be on the same device
x.shape, x.dtype, x.device                        # torch.Size([2, 2]), torch.float32, device(type='cuda')

API differences from NumPy are small but constant: axis is dim, reshape has a sibling view (which requires contiguous memory), np.concatenate is torch.cat, and in-place operations end in an underscore (x.add_(1)). Default float dtype is float32, which is what you want.

Autograd: gradients of any program

Set requires_grad=True on a tensor and PyTorch records every operation applied to it in a directed acyclic graph. Calling .backward() on a scalar output walks that graph in reverse, applying the chain rule (1.4 Derivatives, gradients and the chain rule) node by node — reverse-mode automatic differentiation, the algorithm that 4.4 Backpropagation, fully derived derives — and accumulates $\partial(\text{output})/\partial(\text{leaf})$ into each leaf's .grad. "Accumulates" is literal: gradients add to whatever is already in .grad, so you must zero them before each backward pass.

w = torch.tensor(2.0, requires_grad=True)
x, y = torch.tensor(3.0), torch.tensor(5.0)
loss = (w * x - y) ** 2            # L = (wx − y)²
loss.backward()
w.grad                             # tensor(6.): dL/dw = 2(wx − y)x = 2(6 − 5)(3) = 6

The graph is built dynamically each forward pass, which is why Python control flow (loops, ifs) just works. Three practical consequences: wrap evaluation code in with torch.no_grad(): so no graph is recorded (saves memory and time); call .detach() to get a tensor that is cut out of the graph (for logging a loss value, use loss.item()); and remember that .backward() frees the graph by default, so you cannot call it twice on the same forward pass without retain_graph=True.

Math check. For $L=(wx-y)^2$, $\partial L/\partial w = 2(wx-y)\cdot x$. With $w=2, x=3, y=5$: $2\cdot(6-5)\cdot 3 = 6$. Autograd agrees exactly; it computes the derivative of the program, not a numerical approximation.

The training loop

Every supervised deep-learning program is this loop. Memorize it:

model = torch.nn.Linear(in_features=3, out_features=1).to(device)   # y_hat = x Wᵀ + b
opt = torch.optim.SGD(model.parameters(), lr=0.05)
loss_fn = torch.nn.MSELoss()

for step in range(1000):
    xb, yb = X[idx], Y[idx]            # 1. get a mini-batch (on device)
    pred = model(xb)                   # 2. forward pass
    loss = loss_fn(pred, yb)           # 3. compute the loss (a scalar)
    opt.zero_grad()                    # 4a. clear old gradients
    loss.backward()                    # 4b. backward pass: fills p.grad for every parameter
    opt.step()                         # 5. update parameters: p ← p − lr · p.grad (for SGD)
    if step % 100 == 0:
        print(f"step {step:4d}  loss {loss.item():.4f}")

Step 5 is gradient descent from 1.10 Optimization: from gradient descent to Adam: for each parameter $\theta$, $\theta \leftarrow \theta - \eta\,\partial L/\partial\theta$. Swapping SGD for Adam changes only the update rule (4.5 Optimizers and learning-rate schedules); swapping Linear for a transformer changes only the forward pass; swapping MSELoss for CrossEntropyLoss changes only the objective (4.3 Loss functions catalog). The loop itself is invariant, which is why reading other people's training code gets easy fast.

For linear regression this is exactly the gradient-descent solution of 3.2 Linear regression. For logistic regression (3.3 Logistic regression and softmax classification) change two things: nn.BCEWithLogitsLoss() as the loss (it applies the sigmoid internally in a numerically stable way) and torch.sigmoid(model(xb)) > 0.5 for predictions. For multi-class, nn.Linear(d, C) with nn.CrossEntropyLoss(), which expects raw logits of shape (B, C) and integer targets of shape (B,) — not one-hot vectors.

nn.Module: how models are organized

A model is a class inheriting from nn.Module: parameters and sub-modules are declared in init, the computation in forward. Calling model(x) runs forward plus hooks; model.parameters() yields every learnable tensor (so the optimizer finds them); model.train() and model.eval() toggle layers like dropout and batch norm (4.6 Initialization and normalization, 4.7 Regularization in deep nets) that behave differently during training.

class MLP(torch.nn.Module):
    def __init__(self, d_in, d_hidden, d_out):
        super().__init__()
        self.net = torch.nn.Sequential(torch.nn.Linear(d_in, d_hidden), torch.nn.ReLU(),
                                       torch.nn.Linear(d_hidden, d_out))
    def forward(self, x):
        return self.net(x)

model = MLP(784, 256, 10)
sum(p.numel() for p in model.parameters())        # 784·256 + 256 + 256·10 + 10 = 203,530
torch.save(model.state_dict(), "mlp.pt")          # save weights only (the recommended way)
model.load_state_dict(torch.load("mlp.pt"))       # restore into an identically built model

The parameter count formula is worth internalizing: a linear layer from $d_{in}$ to $d_{out}$ has $d_{in}d_{out}+d_{out}$ parameters. You will count parameters like this for CNNs (5.2 CNN building blocks) and transformers (7.4 Inside the block: residuals, LayerNorm, FFN and embeddings).

Reading shape errors

The two errors you will see most: mat1 and mat2 shapes cannot be multiplied (32x784 and 28x256) means the input width does not match in_features — usually a missing flatten or a transposed weight; and Expected input batch_size (32) to match target batch_size (1) or a loss that is suspiciously constant, which usually means predictions of shape (B, 1) were compared against targets of shape (B,) and broadcasting silently produced a (B, B) matrix. Fix shapes at the boundary (yb.view(-1, 1) or pred.squeeze(-1)) and assert them in forward while developing. Device mismatches (Expected all tensors to be on the same device) mean a tensor was created on CPU inside the loop; create it on device or move it.

Where this shows up later

• The loop is unchanged for CNNs, RNNs and GPT — 4.9 PyTorch II: real training pipelines, 5.4 Transfer learning and training CNNs in practice, 7.8 Build GPT from scratch (nanoGPT walkthrough); only the model, loss and optimizer change.
• What autograd computes is derived by hand in 4.4 Backpropagation, fully derived; Karpathy's micrograd builds a tiny autograd in 100 lines.
• Optimizers beyond SGD are the subject of 4.5 Optimizers and learning-rate schedules; train/eval mode matters for dropout and normalization — 4.6 Initialization and normalization, 4.7 Regularization in deep nets.
• Logits + CrossEntropyLoss is the output head of every classifier and language model — 4.3 Loss functions catalog, 7.6 Decoder-only: GPT and autoregressive generation.