Tensors and autograd

The core data structure, moving work to the GPU, and the automatic differentiation that makes training possible.

Tensors are NumPy with two extras

import torch

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

torch.zeros(2, 3)
torch.ones_like(t)
torch.arange(0, 10, 2)
torch.randn(3, 3)              # standard normal

t + 1
t * t
t @ t                          # matrix multiply
t.mean(dim=0)
t.T

The two extras over NumPy: hardware acceleration (CPU/GPU) and autograd. NumPy and torch tensors convert freely — torch.from_numpy(arr) shares memory.

Autograd

x = torch.tensor(2.0, requires_grad=True)
y = x ** 3 + 2 * x

y.backward()          # compute dy/dx
x.grad                # tensor(14.)  -> 3x² + 2 at x = 2

# during evaluation you do not need the graph
with torch.no_grad():
    preds = model(inputs)

# after a step, clear the accumulated gradients
optimizer.zero_grad()
  • Gradients accumulate until you zero them — forgetting zero_grad() is the classic training bug.
  • Detach tensors (.detach()) when you want a value without tracking history, e.g. logging metrics.
  • Wrap evaluation in torch.no_grad() to save memory and time.
💡
Autograd builds a graph as you compute. Every forward pass that you keep alive holds memory — which is why training loops free it (via loss.backward() then optimizer.step()) rather than keeping many graphs around.

CPU, GPU and dtype

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

t = t.to(device)
model = model.to(device)

# every tensor in an operation must live on the same device
# RuntimeError: Expected all tensors to be on the same device
SettingEffect
float32Default; a good balance of speed and precision
float16 / bfloat16Faster on modern GPUs; use with mixed precision
torch.set_float32_matmul_precision('high')Allows faster matmul on some hardware
.to(device)Copy to CPU/GPU — an error here explains most device mismatches
⚠️
"Expected all tensors to be on the same device" means exactly that: inputs, labels and model must agree. Put the device in one constant and use it everywhere.

FAQ

Is PyTorch eager by default?
Yes — operations execute immediately, which makes debugging with a normal debugger practical. Compilation (torch.compile) can speed up hot loops later.
How is it different from TensorFlow?
Conceptually both do the same job. PyTorch's define-by-run model and Python-native control flow made it the research favourite; the two have converged in capability.

The training loop NumPy arrays

Last refreshed 2026-09-18.