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.TThe 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| Setting | Effect |
|---|---|
float32 | Default; a good balance of speed and precision |
float16 / bfloat16 | Faster 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.
Related
The training loop NumPy arrays
Last refreshed 2026-09-18.