Add a batch dimension
A model expects a batch axis but one sample lacks it.
PyTorch · Function
torch.unsqueeze
Adds a size-one dimension to a tensor at a chosen position. The returned tensor shares the input's underlying data.
How it works
torch.unsqueeze returns a tensor view with one size-one dimension inserted at the requested position. The input values remain available through shared underlying storage, and negative positions are interpreted relative to the resulting rank.
Syntax
torch.unsqueeze(input, dim)Parameters
| Name | Type | Description |
|---|---|---|
input | Tensor | The tensor whose shape receives an added singleton dimension. |
dim | int | The integer position at which to insert the singleton dimension. |
Returns
A new Tensor with one dimension of size one inserted; it shares the input tensor's underlying data.
Key points
- The operation adds exactly one singleton dimension at the requested position.
- The result shares storage with the input, making the operation view-like rather than a value-copying reshape.
- The functional API accepts a tensor and an integer dimension, then returns a Tensor.
- dim identifies the insertion location, not a dimension whose existing values are selected or removed.
Common mistakes
- The dimension index must be within the valid range determined by the input rank; values outside it are invalid.
- A negative dimension is translated using the input rank plus one, so it may not match the equivalent indexing rule for an existing dimension.
Example: Insert a singleton dimension
import torch
x = torch.tensor([1, 2, 3, 4])
y = torch.unsqueeze(x, 1)
assert y.shape == (4, 1)
assert torch.equal(y[:, 0], x)
print(y)Adds a singleton dimension at index 1, converting the four-element vector into a column-shaped tensor.
When to use it
Control broadcasting shape
A vector must align as a row or column in broadcasting or matrix operations.
Prepare tensor dimensions
A scalar or lower-rank tensor needs an explicit axis before concatenation or stacking.
Interview questions
QWhat does torch.unsqueeze do?
It inserts a dimension whose size is one at the position specified by dim. For example, inserting at position 1 changes a one-dimensional tensor of length four into shape (4, 1).
Follow-up: What happens to the tensor's data? The result shares the same underlying data as the input, so this operation changes the view's shape rather than copying values.
QWhat are torch.unsqueeze's parameters and return type?
input is the tensor to reshape conceptually, and dim is the integer index where the singleton dimension is inserted. The function returns a Tensor.
Follow-up: How are negative indices handled? PyTorch maps a negative dim by adding input.dim() + 1.
QWhich dim values are valid?
For an input with rank n, dim may range from -n - 1 through n inclusive, expressed as the half-open interval [-n - 1, n + 1). A valid insertion therefore produces a tensor with rank n + 1.
Follow-up: Why is the upper endpoint exclusive? The documented interval notation permits n but excludes n + 1, which is beyond the available insertion positions.
Related
torch.Tensor.unsqueezetorch.Tensor.unsqueeze_torch.squeeze