Learn AI Create the Future — with Pranay
Main site

Debugging PyTorch

Reading a shape error

Shape errors name two mismatched sizes — learn to factor those numbers back into your batch and feature dimensions, and the bug points at itself.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. Where you have seen this
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

A shape error means two tensors met an operation that needs their sizes to fit together — and they did not.

It is a plug-and-socket problem. A three-pin plug does not enter a two-pin socket, no matter how you push. The plug is fine. The socket is fine. The pairing is wrong.

The error message is unusually helpful, once you know how to read it: it prints the exact pin counts of both sides. Your whole job is to recognise which numbers those are.

Why it exists

Every layer of a model expects input of a promised size and produces output of a promised size. Chain layers together and the promises must line up end to end, like train coaches coupling.

PyTorch checks the coupling at every join. The moment two sizes disagree, it stops and reports both. That strictness is a gift — the alternative would be silently wrong arithmetic.

How it works

your data:  [ 32 photos, each 784 numbers ]
                                │
                    ┌───────────▼───────────┐
                    │ layer expecting 784 ✓ │ → works
                    │ layer expecting 512 ✗ │ → shape error, both numbers printed
                    └───────────────────────┘

Reading the error is detective work with two clues: the two printed sizes. One usually belongs to your data, the other to a layer's expectation. Find which is which, and you have found the bug.

Where you have seen this

Anyone assembling flat-pack furniture has lived this: a bolt that will not thread because it belongs to a different panel. The instructions name the parts; the parts in your hand disagree. You resolve it by checking labels, not by forcing the bolt — same discipline here.

Remember this

  • A shape error prints both mismatched sizes — they are clues, not noise.
  • One number is usually your data, the other a layer's expectation.
  • Factor unfamiliar numbers: a strange size is often two familiar sizes multiplied.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Captured with torch 2.5.1 on CPU. Error wording can shift slightly between versions; the structure stays.

A classic crash, on purpose

Images into a linear layer, with the flattening forgotten:

shape_crash.py
import torch
import torch.nn as nn

torch.manual_seed(0)
images = torch.randn(32, 1, 28, 28)          # a batch of 32 greyscale 28x28 images
model = nn.Linear(784, 10)                   # expects flat vectors of 784 numbers
out = model(images)
Output
RuntimeError: mat1 and mat2 shapes cannot be multiplied (896x28 and 784x10)

Decoding it, step by step

mat1 is your data, mat2 is the layer's weights. For a Linear(784, 10), mat2 being 784x10 confirms it: 784 in, 10 out.

Where did 896x28 come from? Your input was (32, 1, 28, 28) — no 896 anywhere. Factor it: 896 = 32 × 1 × 28. A Linear layer treats only the last dimension as features and folds every other dimension into rows. So the 4-D image batch became 896 rows of 28 features, and 28 does not match the expected 784.

That factoring move — decomposing an unfamiliar number into your known sizes — cracks most shape errors in seconds. Seeing 896, you should think "32 times 28", and the missing flatten announces itself.

The fix

shape_fixed.py
import torch
import torch.nn as nn

torch.manual_seed(0)
images = torch.randn(32, 1, 28, 28)
print("before:", images.shape)

flat = images.flatten(start_dim=1)           # keep dim 0 (the batch), flatten the rest
print("after: ", flat.shape)

model = nn.Linear(784, 10)
out = model(flat)
print("output:", out.shape)
Output
before: torch.Size([32, 1, 28, 28])
after:  torch.Size([32, 784])
output: torch.Size([32, 10])

flatten(start_dim=1) preserves the batch dimension and merges the rest: 1 × 28 × 28 = 784. Inside an nn.Sequential, the same job is done by an nn.Flatten() layer.

The general triage method

  1. Read both shapes in the message. Write them down.
  2. Print shapes at the boundaries. A print(x.shape) between stages is not shameful; it is the fastest instrument you own.
  3. Factor strange numbers into products of your known sizes: batch, channels, height, width, sequence length.
  4. Check the convention. Linear wants (..., features); Conv2d wants (N, C, H, W); CrossEntropyLoss wants logits (N, C) against targets (N,). Half of all shape errors are two conventions meeting.

Common mistakes

Fixing the error by matching numbers blindly. Changing Linear(784, 10) to Linear(28, 10) silences this crash — and trains a model where every image row is a separate sample. The error disappears; the bug becomes invisible. Understand which side is wrong before editing either.

view instead of flatten on non-contiguous tensors. x.view(32, -1) can raise its own error after transposes or slicing. flatten/reshape handle those cases — the full story is in reshape, view and contiguity.

Flattening the batch dimension too. x.flatten() with no arguments produces one long vector of 25088 numbers, and the next error reports sizes that confuse you further. Always start_dim=1.

Trusting that no error means right shapes. Broadcasting can make mismatched tensors combine without complaint — a (32, 1) meeting a (32,) produces a silent (32, 32). Shape errors are the friendly failure; the silent ones cost more.

Try it yourself

Feed the fixed model a batch of (32, 3, 28, 28) colour images and predict the exact numbers the error will print, before running it. If your prediction matches, the triage method is yours.

What to learn next

Researcher — Mathematics and papers.

The rule being enforced

nn.Linear computes $y = x W^\top + b$ with $W \in \mathbb{R}^{d_{out} \times d_{in}}$, accepting input $x$ of shape $(, d_{in})$ where $$ is any number of leading dimensions; internally the leading dimensions are folded, the matmul runs as $(\prod *, d_{in}) \times (d_{in}, d_{out})$, and the result unfolds. The error's 896x28 is that folded view leaking into the message — a useful implementation detail to recognise.

For batched matmul, shapes $(\ldots, n, m) @ (\ldots, m, p)$ require the trailing pair to agree while leading dimensions broadcast per the standard rules; errors there report either the trailing mismatch or a broadcast failure, two different messages worth telling apart.

Why shapes cannot be checked earlier

Python-level PyTorch is dynamically shaped: shapes are values, known only when tensors exist. Static verification is an active research and tooling area:

  • Named tensors — Rush's Tensor Considered Harmful (2019) argued for dimension names over positions; PyTorch's named-tensor prototype has since been deprecated, but the critique shaped later tools.
  • jaxtyping-style annotations — Float[Tensor, "batch 784"] — give runtime-checked shape contracts and, with a checker, documentation that cannot rot.
  • torch.export / symbolic shapes — the compile stack traces shapes as symbols (s0, s1) with guards, which is why torch.compile errors sometimes speak of guard failures rather than concrete sizes: the same mismatch, one abstraction level up.

Error-message archaeology

The same underlying mismatch surfaces differently by operation: mat1 and mat2 shapes cannot be multiplied (linear/matmul), expected input[32, 28, 28] to have 1 channels, but got 28 channels (conv layers, which check (N, C, H, W) semantically), The size of tensor a (X) must match the size of tensor b (Y) at non-singleton dimension d (broadcast failures), and Expected input batch_size (X) to match target batch_size (Y) (losses). Mapping message-family to operation-family is most of the diagnostic skill this lesson trains.

What to learn next