PyTorch in 5 minutes

PyTorch is NumPy with two extras: it runs on GPUs, and it tracks gradients automatically via autograd (a built-in engine that records every operation on a tensor so it can later compute derivatives). That’s the entire pitch. Everything you’ve learned about array shapes, broadcasting, and indexing carries straight over.

Tensors are arrays with a backstory

A tensor is a multi-dimensional array — same as an ndarray — but it remembers operations performed on it (so it can compute gradients later) and it can live on a GPU.

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
print(x.shape)        # torch.Size([2, 2])
print(x.dtype)        # torch.float32
print(x.device)       # cpu

The three attributes you’ll check constantly: shape, dtype, device. Most “RuntimeError” messages in PyTorch trace back to a mismatch on one of these three.

The training loop, in five steps

The shape of every PyTorch training step is identical:

pred = model(x)              # 1. forward pass: run inputs through the model
loss = loss_fn(pred, y)      # 2. compute the loss (how wrong we are)
loss.backward()              # 3. backward pass: autograd fills in .grad for every parameter
optimizer.step()             # 4. update weights using those gradients
optimizer.zero_grad()        # 5. clear gradients so they don't accumulate into next step

The forward pass runs data through the model to get a prediction. The backward pass walks the recorded operations in reverse to compute how much each weight contributed to the loss (the gradient). Step 5 is critical: PyTorch adds to .grad rather than replacing it, so without zero_grad() each step’s gradients stack up and the updates are wrong.

That’s it. Whether you’re training a 100-parameter regression or a 70B language model, this loop is the heartbeat. The rest is data loading, logging, and checkpoints.

Linear regression, the PyTorch way

Let’s fit y = 2x + 1 with noise — the smallest possible training run. PyTorch can’t run in the browser, so here is the same loop in NumPy — with its real output — then the identical PyTorch version beside it.

import numpy as np

# Generate noisy data: y = 2x + 1 + noise
rng = np.random.default_rng(0)
x = rng.uniform(-2, 2, size=(64, 1)).astype(np.float32)
y = 2.0 * x + 1.0 + 0.1 * rng.standard_normal(x.shape).astype(np.float32)

# Initialize weights randomly (the "model").
w = rng.standard_normal((1, 1)).astype(np.float32)
b = np.zeros((1,), dtype=np.float32)

lr = 0.1
for step in range(200):
    # 1. Forward: prediction = x @ w + b
    pred = x @ w + b
    # 2. Loss: mean squared error
    loss = ((pred - y) ** 2).mean()
    # 3. Backward (manual): d_loss / d_w, d_loss / d_b
    grad_pred = 2 * (pred - y) / x.shape[0]
    grad_w = x.T @ grad_pred
    grad_b = grad_pred.sum(axis=0)
    # 4. Update
    w -= lr * grad_w
    b -= lr * grad_b

print(f"learned w = {w.item():.3f}, b = {b.item():.3f}")
print(f"target  w = 2.000, b = 1.000")
print(f"final loss = {loss:.5f}")

learned w = 2.008, b = 1.006
target  w = 2.000, b = 1.000
final loss = 0.00968

Notice we computed gradients by hand. The whole point of PyTorch is to not have to do that. Here’s the same training loop in real PyTorch:

import torch
import torch.nn as nn

x = torch.randn(64, 1)
y = 2.0 * x + 1.0 + 0.1 * torch.randn(64, 1)

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

for step in range(200):
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

print(model.weight.item(), model.bias.item())

Device, the one-line GPU switch

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
x = x.to(device)
y = y.to(device)

Tensors and models have to be on the same device. Forgetting .to() on one of them is the #2 source of PyTorch bugs (the #1 is shape mismatches). On Apple Silicon, use "mps" instead of "cuda".

Why this matters in industry

Every modern deep learning framework — JAX, TensorFlow, MLX — converged on this same forward/backward/step structure because it composes cleanly. Once you know the PyTorch loop, you can read any of the others. PyTorch is the default in research labs and ~80% of production model training today.

In one breath

• A PyTorch tensor is a NumPy array with two extras: it runs on GPUs, and it records its operations so autograd can differentiate them.
• Check shape, dtype, and device constantly — most PyTorch RuntimeErrors are a mismatch in one of those three.
• Every training step is the same five lines: forward → loss → backward → step → zero_grad.
• zero_grad is not optional: backward adds to .grad, so without it gradients stack across steps and the update is wrong.
• Move model and data to the same device with .to(device); the loop is identical from a 100-parameter fit to a 70B model.

Quick check