Introduction

PyTorch models care about tensor shape just as much as tensor values. squeeze and unsqueeze are small operations, but they solve a huge number of shape mismatches by either removing dimensions of size 1 or inserting new singleton dimensions exactly where a model expects them.

What squeeze Does

torch.squeeze removes dimensions whose length is 1. That is useful when data has extra wrapper dimensions that do not carry real information.

1import torch
2
3x = torch.randn(1, 3, 1, 5)
4print(x.shape)
5
6y = x.squeeze()
7print(y.shape)

Output:

torch.Size([1, 3, 1, 5])
torch.Size([3, 5])

Without an explicit dimension argument, squeeze removes every singleton dimension. That is convenient, but it can also remove a batch dimension you wanted to keep.

You can be more precise:

1import torch
2
3x = torch.randn(1, 3, 1, 5)
4y = x.squeeze(2)
5
6print(y.shape)

Now only dimension 2 is removed, so the batch dimension remains.

What unsqueeze Does

torch.unsqueeze does the opposite. It inserts a dimension of size 1 at the given position.

1import torch
2
3x = torch.tensor([1.0, 2.0, 3.0])
4print(x.shape)
5
6y = x.unsqueeze(0)
7z = x.unsqueeze(1)
8
9print(y.shape)
10print(z.shape)

Output:

torch.Size([3])
torch.Size([1, 3])
torch.Size([3, 1])

This is common when you have a single sample and need to add a batch axis before passing it to a model.

Typical Deep Learning Use Cases

The most common pattern is preparing tensors for layers that expect a fixed rank.

For example, a grayscale image loaded as height by width may need a channel dimension and a batch dimension:

1import torch
2
3image = torch.randn(28, 28)
4image = image.unsqueeze(0).unsqueeze(0)
5
6print(image.shape)

The result is shape [1, 1, 28, 28], which matches the usual batch, channel, height, width convention used by convolutional layers.

The reverse happens after model output or indexing. You may end up with a tensor such as [1, 10] and want to remove the batch axis before post-processing:

1logits = torch.randn(1, 10)
2scores = logits.squeeze(0)
3
4print(scores.shape)

Dimension Choice Matters

Both functions are simple, but the dimension index controls semantics. Inserting a singleton dimension at 0 means "new batch axis" in many pipelines. Inserting at 1 may mean "new channel axis." Removing the wrong singleton dimension can subtly break downstream layers even when the tensor still has a valid shape.

That is why many experienced PyTorch users prefer explicit calls such as unsqueeze(0) or squeeze(1) instead of relying on the no-argument version everywhere.

In-Place Variants

PyTorch also provides squeeze_ and unsqueeze_, which modify the tensor in place.

1import torch
2
3x = torch.randn(1, 5)
4x.squeeze_(0)
5
6print(x.shape)

In-place shape operations are fine in some preprocessing code, but they make debugging harder if multiple variables reference the same tensor. Unless you need the in-place behavior, the non-mutating form is usually easier to reason about.

Common Pitfalls

• Calling squeeze() without a dimension can accidentally remove the batch axis when batch size is 1.
• 'squeeze(dim) removes a dimension only if its size is 1. If the chosen dimension is larger, nothing happens.'
• Using the wrong unsqueeze position can swap the meaning of batch and channel dimensions.
• Expecting squeeze or unsqueeze to reorder data is a mistake. They change shape metadata, not element order.
• Chaining many shape operations without printing tensor.shape makes shape bugs much harder to diagnose.

Summary

• 'squeeze removes dimensions of size 1, while unsqueeze adds a size 1 dimension at a chosen index.'
• Use unsqueeze to add batch or channel axes before model input.
• Use squeeze to remove wrapper dimensions after indexing or inference output.
• Prefer explicit dimension arguments when the meaning of a shape matters.