PyTorch unsqueeze Explained | Generated by AI

Home 2026.07

Question: What does unsqueeze do in PyTorch?

Answer: Inserts a dimension of size 1 at the specified position. It’s the inverse of squeeze. No data copy, just a view (metadata change on strides).

import torch

x = torch.arange(6).reshape(2, 3)  # shape [2, 3]

x.unsqueeze(0).shape   # [1, 2, 3]  -> insert at front
x.unsqueeze(1).shape   # [2, 1, 3]  -> insert in middle
x.unsqueeze(-1).shape  # [2, 3, 1]  -> insert at end
x.unsqueeze(2).shape   # [2, 3, 1]  -> same as -1 here

Index semantics: for tensor of dim n, valid range for dim is [-(n+1), n]. The new dim gets inserted before that index in the output shape.

Common use case — batch dim for a single sample:

img = torch.randn(3, 224, 224)      # C, H, W
batch = img.unsqueeze(0)            # [1, 3, 224, 224]

Or broadcasting alignment, e.g. adding a bias per-row:

a = torch.randn(4, 5)
b = torch.randn(4)
c = a + b.unsqueeze(1)  # b: [4] -> [4,1], broadcasts against [4,5]

It’s a view op — x.unsqueeze(0).data_ptr() == x.data_ptr() is True. Gradients flow through fine since it’s just a reshape in the autograd graph.

References:


Generated by AI. Curating and sharing still takes effort. If you find it useful, feel free to donate. WeChat: @lzwjavaWeChat QR · X: @lzwjava · Say hi 👋

Back Donate