What Does `unsqueeze` Do in PyTorch? (And Why Your Model Keeps Asking For It)

You passed a single image to your model and got ValueError: expected 4D input (got 3D input). Someone on Stack Overflow said "just add .unsqueeze(0)" — and it worked. But nobody explained what it did, or how you'd know to reach for it next time.

The short answer

unsqueeze(dim) inserts a new axis of size 1 at position dim. The data doesn't change at all — only the shape gains an extra 1. A tensor of shape (4,) becomes (1, 4) with unsqueeze(0) or (4, 1) with unsqueeze(1). Its inverse, squeeze(dim), removes a size-1 axis. Both are free: no memory is allocated and no data is copied.

import torch

v = torch.tensor([1, 2, 3, 4])   # shape (4,)

row = v.unsqueeze(0)             # (1, 4) — one row
col = v.unsqueeze(1)             # (4, 1) — one column

print(row.shape, col.shape)
# torch.Size([1, 4]) torch.Size([4, 1])

Same four numbers, three different shapes. That's the whole idea.

Why does PyTorch need size-1 dimensions at all?

Because a shape isn't just a size — it's a meaning. In PyTorch, shape (4,) means "four numbers." Shape (1, 4) means "one row of four." Shape (4, 1) means "four rows of one." Those are three different claims about your data, and PyTorch takes them literally.

That's why so many frustrating shape errors happen not because your data is the wrong size, but because it's missing a dimension — or has one too many. A model expects a batch of inputs and you hand it a single sample. A loss function wants a column and you give it a flat row. The fix, nearly every time, is one call.

The number you pass is where the new axis goes. Dimension 0 is the front, dimension 1 is next, and negative indexing works too — unsqueeze(-1) always appends a new last axis no matter how many dimensions the tensor already has.

t = torch.randn(3, 4)

print(t.unsqueeze(0).shape)    # torch.Size([1, 3, 4])
print(t.unsqueeze(1).shape)    # torch.Size([3, 1, 4])
print(t.unsqueeze(-1).shape)   # torch.Size([3, 4, 1])

The None shortcut (same thing, fewer characters)

There's a terser spelling you'll see constantly in real codebases: put None inside an index expression. Wherever you place None, a new size-1 axis appears.

v = torch.tensor([1, 2, 3, 4])

row = v[None, :]     # same as unsqueeze(0)
col = v[:, None]     # same as unsqueeze(1)

print(row.shape, col.shape)
# torch.Size([1, 4]) torch.Size([4, 1])

Both spellings do exactly the same thing. None indexing is shorter; unsqueeze is more explicit about where the axis lands. Use whichever reads more clearly — you'll meet both in the wild.

What does squeeze do in PyTorch?

squeeze is the inverse: it removes dimensions whose size is 1. With no argument it removes all of them; with a dim argument it removes only that one — and only if it really is size 1.

t = torch.randn(1, 3, 1, 4)

print(t.shape)                # torch.Size([1, 3, 1, 4])
print(t.squeeze().shape)      # torch.Size([3, 4])     — all 1s gone
print(t.squeeze(0).shape)     # torch.Size([3, 1, 4])  — only dim 0
print(t.squeeze(1).shape)     # torch.Size([1, 3, 1, 4]) — dim 1 is size 3, unchanged

Look at that last line. Squeezing dimension 1 did nothing, because dimension 1 has size 3, not 1. squeeze only removes axes that are actually size 1, so it's safe to call even when you're not sure.

Safe, but not always wise — see the gotchas below.

A worked example: the batch dimension

This is the case that sends most beginners to Google, so let's walk it end to end.

Most PyTorch models expect a batch — shape (N, ...), where N is how many samples you're passing. A single CIFAR-style image is shape (3, 32, 32): three channels, 32 by 32 pixels. That's three dimensions, and a lot of layers want four.

Here's the part that trips people up, because it's changed over the years: a bare nn.Conv2d will actually accept your unbatched image just fine. Modern PyTorch added no-batch-dim support to many nn modules, so this runs:

import torch
import torch.nn as nn

conv = nn.Conv2d(3, 16, kernel_size=3)
sample = torch.randn(3, 32, 32)      # one image — 3 dims

print(conv(sample).shape)            # torch.Size([16, 30, 30]) — no error

So why does everyone tell you to unsqueeze? Because the moment there's a normalization layer in the stack, the tolerance ends. BatchNorm2d computes statistics across the batch, so it genuinely cannot work without a batch axis:

model = nn.Sequential(
    nn.Conv2d(3, 16, kernel_size=3),
    nn.BatchNorm2d(16),
)

model(sample)

ValueError: expected 4D input (got 3D input)

Read it literally: it got 3 dimensions and wanted 4. Nothing is wrong with your image — it's just not in a batch. unsqueeze(0) wraps it in a batch of exactly one:

batch = sample.unsqueeze(0)          # (1, 3, 32, 32)
output = model(batch)

print(batch.shape)                   # torch.Size([1, 3, 32, 32])
print(output.shape)                  # torch.Size([1, 16, 30, 30])

The lesson generalizes: whether a single layer tolerates a missing batch dimension is a detail that varies by layer and by version. Whether your model does is not worth gambling on. Add the batch axis and the question disappears.

And on the way out, you often want that batch axis gone again, or a trailing singleton removed — a regression head that outputs (32, 1) won't line up with targets shaped (32,):

preds = torch.randn(32, 1)           # model output
print(preds.squeeze(-1).shape)       # torch.Size([32])

unsqueeze going in, squeeze coming out. That pairing shows up in almost every inference script you'll ever write.

The other big use: lining up shapes for broadcasting

The second place unsqueeze earns its keep is broadcasting. To subtract a per-channel mean from an image, the mean has to be shaped so it aligns with the channel axis — not the width axis.

mean = torch.tensor([0.485, 0.456, 0.406])   # (3,)
img = torch.randn(3, 224, 224)               # (C, H, W)

img - mean

RuntimeError: The size of tensor a (224) must match the size of tensor b (3)
at non-singleton dimension 2

Broadcasting aligns shapes from the right, so (3,) tries to line up with the 224-wide last axis and fails. Give it two trailing size-1 axes and it lines up correctly:

mean = mean[:, None, None]           # (3, 1, 1)
img = img - mean                     # broadcasts to (3, 224, 224)

print(mean.shape, img.shape)
# torch.Size([3, 1, 1]) torch.Size([3, 224, 224])

The 1s are placeholders that say "stretch me along this axis." That's what makes the subtraction apply per channel instead of per pixel column. If the right-to-left alignment rule is new to you, I wrote up how PyTorch broadcasting works in full here.

Are squeeze and unsqueeze expensive?

No — they're free. Both return a view of the same data with adjusted strides. No memory is allocated, nothing is copied, so you can call them freely, even inside a hot loop.

v = torch.arange(4)
u = v.unsqueeze(0)

print(u.data_ptr() == v.data_ptr())   # True
print(u.stride())                     # (4, 1)

Identical data pointer means no copy happened. unsqueeze just added a stride entry. This is the same principle behind view and reshape: any shape change expressible as a stride change is a view; anything else needs a copy.

Common mistakes and gotchas

• Bare squeeze() eating your batch dimension. If your batch size happens to be 1, squeeze() with no argument silently deletes it along with every other singleton — and your shapes break several lines later, far from the cause. Prefer squeeze(dim), always.
• Confusing unsqueeze(1) with unsqueeze(-1). They're the same only for 1-D tensors. On a (3, 4) tensor, unsqueeze(1) gives (3, 1, 4) while unsqueeze(-1) gives (3, 4, 1).
• Expecting squeeze(dim) to raise on a non-1 axis. It doesn't — it quietly returns the tensor unchanged. Convenient, but it means a typo'd dim fails silently.
• Reaching for reshape where unsqueeze is clearer. x.reshape(1, *x.shape) works, but x.unsqueeze(0) says what you mean.
• Forgetting the batch axis at inference time. Training works because your DataLoader adds the batch dimension for you. Then you feed a single sample by hand and it breaks. That's not a new bug — it's the DataLoader no longer doing you a favor.
• Trusting that "it ran" means the shape was right. Since some nn modules accept unbatched input, a missing batch dimension can survive several layers before something downstream complains — or worse, quietly produce output of the wrong rank. Print .shape.

Recap

unsqueeze(dim) inserts a size-1 axis at dim; squeeze(dim) removes one. v[None, :] is the same thing written inline. Both are free views — no copy. Use unsqueeze(0) to add a batch dimension for a single sample, unsqueeze(-1) or [:, None, None] to align shapes for broadcasting, and squeeze(-1) to drop a trailing singleton from model outputs.

---

This is one idea from my book PyTorch From Ground Up, which builds everything from tensors upward so nothing stays vague. If it helped, you can grab the free PyTorch tensor cheat-sheet here, run every example from the book in the companion notebooks on GitHub, or find the full paperback on Amazon here.