Search Tech Journey

Find topics, journeys and posts

6-month learning plan100 / 130
back to blog
mladvanced 15m read

S100 · PyTorch Fundamentals — Tensors, Autograd, nn.Module

The tool you'll live in.

Module M12: Deep Learning · Session 100 of 130 · Track: ML 🔥

What you'll be able to do after this session

  • Explain PyTorch Fundamentals to a friend in one minute, out loud, no notes.
  • Recognise it when you see it in real code / systems / papers.
  • Complete the hands-on exercise at the end without looking up help.

Prerequisites

If you can't answer these in 30 seconds each, redo the linked session first:


Pre-read (skim before the session)


(a) Intuition · 5 min

The one-paragraph mental model. Why this concept exists, what problem it solves, and the analogy that makes it stick.

In one sentence: The tool you'll live in.

Imagine you're doing your taxes on a giant spreadsheet. Every cell has a formula that depends on other cells. When you change one number at the top, Excel walks the dependency graph and updates every downstream cell — you never re-derive the chain rule by hand. PyTorch is Excel for math on GPUs. A Tensor is a cell (well, a multi-dimensional cell — a number, a vector, a matrix, or a 4D block of pixels). Every operation on a tensor is a formula. When you finally call .backward() on a loss, PyTorch walks the graph backwards and hands you the derivative of the loss with respect to every parameter — that's autograd. Then nn.Module is just a folder you keep your parameters in so you don't have to pass 400 weight tensors around by hand.

PyTorch exists because before it (2016 and earlier), deep learning frameworks — Theano, Caffe, early TensorFlow — asked you to first declare a static computation graph, then feed data through it. Debugging was miserable: a shape mismatch showed up 50 layers away from where you wrote the bug. PyTorch went the opposite way: define-by-run. The graph is built on the fly as your Python code runs, so a stack trace points at the exact line where things went wrong. You can print(tensor.shape) in the middle of a forward pass. This one design choice is why researchers switched, why HuggingFace is written in PyTorch, and why almost every 2023+ open model ships PyTorch weights first.

Gotcha to carry forever: a tensor only tracks gradients if requires_grad=True (parameters do this by default via nn.Parameter; raw torch.tensor([1.,2.]) does not). And .backward() accumulates gradients into .grad — you must call optimizer.zero_grad() (or param.grad = None) before each step, or your gradients silently add up across batches and training explodes.


(b) Visual walkthrough · 15 min

Diagrams, worked examples, step-by-step visuals.

PyTorch has three layers you must hold in your head at once. Data lives in Tensors. Operations on tensors get recorded into a dynamic computation graph so autograd can differentiate. Parameters + forward logic get bundled inside an nn.Module. Here's how they compose during one training step:

Worked example — a single autograd trace. Suppose x = torch.tensor(2.0, requires_grad=True) and we compute y = x**2 + 3*x. That's y = 4 + 6 = 10. Behind the scenes PyTorch built this graph: x → Pow(2) → 4, x → Mul(3) → 6, add → 10. When you call y.backward(), PyTorch applies the chain rule from y back to x: dy/dx = 2x + 3 = 2·2 + 3 = 7. Read it off x.grad and it says tensor(7.). You never wrote a derivative — you wrote the forward pass, PyTorch remembered the ops, then differentiated them.

Second worked example — tensor shape arithmetic. A batch of 32 MNIST images has shape (32, 1, 28, 28). You flatten with x.view(32, -1)(32, 784). Pass through nn.Linear(784, 128)(32, 128). ReLU keeps shape. Then nn.Linear(128, 10)(32, 10) — one logit per digit class per image. Cross-entropy takes (32, 10) logits + (32,) integer labels → single scalar loss. The scalar is critical: .backward() only works on a scalar (or you must pass a gradient= argument). This shape choreography is 80% of debugging PyTorch code — always print(x.shape) when confused.

ConceptAPICommon mistake
Create a leaf that trainstorch.randn(3, requires_grad=True) or nn.Parameter(...)Forgetting requires_grad; then .grad is None forever.
Move to GPUx = x.to('cuda')Model on GPU, data on CPU → RuntimeError.
Detach from graphx.detach() or with torch.no_grad():Doing eval inside the graph → OOM on huge models.
Reshape.view() (contiguous) vs .reshape() (safe).view() after .transpose() crashes; use .reshape().

(c) Hands-on · 20 min

Code you type and run. Not pseudo-code — real, runnable, copy-pasteable.

Install once: pip install torch torchvision. Then save as pt_basics.py and run python pt_basics.py.

import torch
import torch.nn as nn
import torch.nn.functional as F
 
torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[info] running on {device}")
 
# 1) TENSORS — the NumPy you always wanted, but on a GPU with gradients.
a = torch.tensor([[1., 2.], [3., 4.]], device=device)
b = torch.randn(2, 2, device=device)
print("a @ b =\n", a @ b)          # matrix multiply
print("shapes:", a.shape, b.shape)  # (2,2) (2,2)
 
# 2) AUTOGRAD — derivatives for free.
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x        # y = 10
y.backward()              # dy/dx = 2x + 3 = 7
print("dy/dx at x=2:", x.grad.item())
 
# 3) nn.Module — a tiny MLP that classifies 10 fake digit classes.
class TinyMLP(nn.Module):
    def __init__(self, in_dim=784, hidden=128, out_dim=10):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden)
        self.fc2 = nn.Linear(hidden, out_dim)
 
    def forward(self, x):
        return self.fc2(F.relu(self.fc1(x)))
 
model = TinyMLP().to(device)
print(model)
print("param count:", sum(p.numel() for p in model.parameters()))
 
# 4) ONE TRAINING STEP on fake data.
X = torch.randn(32, 784, device=device)          # batch of 32 "images"
targets = torch.randint(0, 10, (32,), device=device)
 
loss_fn = nn.CrossEntropyLoss()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
 
for step in range(5):
    opt.zero_grad()          # <-- MUST zero, else grads accumulate
    logits = model(X)        # forward, builds graph
    loss = loss_fn(logits, targets)
    loss.backward()          # autograd fills every .grad
    opt.step()               # W -= lr * W.grad
    print(f"step {step}: loss={loss.item():.4f}")
 
# 5) INFERENCE — turn autograd OFF to save memory + go faster.
model.eval()
with torch.no_grad():
    preds = model(X).argmax(dim=1)
print("first 8 preds:", preds[:8].tolist())

What to observe (checklist):

  1. The [info] running on ... line — if it says cpu but you have a GPU, your CUDA install is broken.
  2. dy/dx at x=2: 7.0 — proves autograd works.
  3. Loss goes down across the 5 steps. If it goes up or NaNs, learning rate is too high.
  4. param count for TinyMLP is 784*128 + 128 + 128*10 + 10 = 101770. Count matches? Good.
  5. preds are integers in [0, 9] — shape and dtype are correct for classification.

Try this modification: comment out opt.zero_grad(). Watch the loss curve go crazy within 5 steps. That single line is the most common bug in beginner PyTorch code — feel the pain once so you never forget it.


(d) Production reality · 10 min

How this shows up in real systems, gotchas, war stories, common bugs.

The four bugs that will bite you in month one. (1) Forgetting .to(device) on either model or batch. You'll get RuntimeError: Expected all tensors to be on the same device. Fix: at the top of your training loop, X, y = X.to(device), y.to(device), and once at model creation, model.to(device). (2) Forgetting optimizer.zero_grad(). Gradients accumulate across steps and your loss diverges after ~10 iterations. Some frameworks (HuggingFace Trainer) hide this — when you drop to raw PyTorch, remember it. (3) Calling .backward() twice on the same graph. By default PyTorch frees the graph after backward to save memory; a second call throws. If you truly need it, pass retain_graph=True — but usually you're just misunderstanding your loop. (4) Doing eval without torch.no_grad(). Building the graph during eval wastes GPU memory and can OOM a model that trains fine.

How real teams ship it. Meta trains LLaMA on PyTorch with FSDP (Fully Sharded Data Parallel) — the model doesn't even fit on one GPU, so shards are moved in and out of each GPU's HBM during each layer's forward/backward. OpenAI, Anthropic, and Mistral all use PyTorch + custom CUDA kernels. HuggingFace's Trainer and PyTorch Lightning are the two dominant "boilerplate remover" wrappers, but in interviews you'll be asked raw PyTorch — always learn the substrate first. For inference at scale, teams compile the model: torch.compile(model) (PyTorch 2.x) can give 30–200% speedup for free by fusing ops via TorchInductor. For real production serving, export to ONNX or TorchScript, or serve directly with vLLM / TensorRT-LLM — raw PyTorch python inference is ~5–10× slower than a compiled runtime and shouldn't be your final deployment path.

One war story: at any medium-sized ML shop, someone will one day report "training loss is fine, validation is garbage". 40% of the time the cause is: model.eval() was never called before validation, so Dropout was still active and BatchNorm was still updating running stats from the val set. Add model.train() / model.eval() to your loop from day one — muscle memory, non-negotiable.


(e) Quiz + exercise · 10 min

  1. What does requires_grad=True do to a tensor, and what happens to .grad if you forget it?
  2. Why must you call optimizer.zero_grad() before every training step? Describe the failure mode if you skip it.
  3. Explain the difference between .view() and .reshape(). When does .view() throw an error?
  4. What is the purpose of with torch.no_grad(): during inference? Name two concrete benefits.
  5. Which method — model.train() or model.eval() — must be called before validation, and which two layer types behave differently in the two modes?

Stretch problem (connects to S098 — Backpropagation): in S098 you derived backprop by hand for a 2-layer net. Re-implement that same 2-layer net in PyTorch using only torch.Tensor operations (no nn.Linear, no nn.Module) with requires_grad=True on your weight tensors, then verify that loss.backward() gives you exactly the same gradients you computed on paper (to within 1e-6). This proves autograd is not magic — it's just chain rule executed by a graph walker.


Explain-out-loud test

If you can't teach these 3 things to a friend without notes, redo the session:

  1. What is PyTorch Fundamentals? (one sentence, no jargon)
  2. When would you use it? (one concrete example)
  3. When would you NOT use it? (one anti-pattern)

What comes next


Part of a 130-session evergreen learning series. Session structure: (a) intuition · (b) visual walkthrough · (c) hands-on · (d) production reality · (e) quiz + exercise. Duration: 60 minutes.

More from M12 · Deep Learning

all modules →
  1. 096Perceptron & Activation Functions
  2. 097Multi-Layer Perceptron — Forward Pass
  3. 098Backpropagation — Derived by Hand on a 2-Layer Net
  4. 099Optimizers — SGD, Momentum, Adam, RMSprop
  5. 101Regularization in DL — Dropout, BatchNorm, Weight Decay
  6. 102CNNs — Convolution, Pooling, ImageNet Architectures