tensor conversion
N-dimensional tensor
tensor manipulation
neural network preprocessing
PyTorch tensor operations

How to convert a list of tensors of dim N to a tensor of dim N1

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

A common preprocessing step in PyTorch is turning a Python list of tensors into one larger tensor with an extra leading dimension. In practice, this usually means taking several tensors that all have shape like [C, H, W] and combining them into a batch with shape [B, C, H, W].

Use torch.stack When Shapes Match

If every tensor in the list has the same shape, the correct tool is usually torch.stack. It creates a new dimension and places each input tensor along that dimension.

python
1import torch
2
3samples = [
4    torch.tensor([1.0, 2.0, 3.0]),
5    torch.tensor([4.0, 5.0, 6.0]),
6    torch.tensor([7.0, 8.0, 9.0]),
7]
8
9batched = torch.stack(samples)
10
11print(batched)
12print(batched.shape)

Output:

python
1tensor([[1., 2., 3.],
2        [4., 5., 6.],
3        [7., 8., 9.]])
4torch.Size([3, 3])

Each original tensor was one-dimensional, so stacking produced a two-dimensional tensor. More generally, stacking tensors of dimension N produces a tensor of dimension N + 1.

You can also choose where the new dimension goes.

python
1images = [torch.randn(3, 32, 32) for _ in range(4)]
2
3batch_first = torch.stack(images, dim=0)
4channel_first_group = torch.stack(images, dim=1)
5
6print(batch_first.shape)
7print(channel_first_group.shape)

The first result has shape [4, 3, 32, 32]. The second inserts the new dimension after the first existing one, producing [3, 4, 32, 32].

Why torch.cat Is Different

It is easy to confuse torch.stack with torch.cat. Concatenation joins tensors along an existing dimension. Stacking adds a new dimension first.

python
1import torch
2
3a = torch.tensor([1, 2, 3])
4b = torch.tensor([4, 5, 6])
5
6print(torch.cat([a, b]).shape)
7print(torch.stack([a, b]).shape)

The first shape is [6], because the vectors were joined end to end. The second shape is [2, 3], because a new leading dimension was created.

That distinction matters in training code. A model expecting batched inputs almost always wants stacking, not concatenation.

Handling Existing Singleton Dimensions

Sometimes the list already contains tensors with an extra singleton axis. In that case, you may be able to remove it with squeeze before stacking, or add it explicitly with unsqueeze if it is missing.

python
1import torch
2
3x = torch.tensor([1.0, 2.0, 3.0])
4print(x.shape)
5
6x2 = x.unsqueeze(0)
7print(x2.shape)

unsqueeze is useful when you have a single tensor and want to simulate a batch of size one. It is not a replacement for stacking multiple tensors, but the concepts are related: both operations increase dimensionality by one.

Variable-Length Inputs Need Padding

torch.stack requires every tensor to have the same shape. If one tensor is shorter or has a different image size, PyTorch raises an error. In that case, you must first normalize the shapes by padding, cropping, or resizing.

python
1from torch.nn.utils.rnn import pad_sequence
2import torch
3
4seqs = [
5    torch.tensor([1, 2, 3]),
6    torch.tensor([4, 5]),
7    torch.tensor([6])
8]
9
10padded = pad_sequence(seqs, batch_first=True, padding_value=0)
11print(padded)
12print(padded.shape)

This is common in natural language and time-series work, where examples do not naturally have uniform lengths.

Common Pitfalls

The most common mistake is using torch.tensor(list_of_tensors). That can trigger an unnecessary copy or fail unexpectedly depending on the tensor contents and device placement. Prefer torch.stack when the inputs are already tensors.

Another issue is mixing devices or data types. If some tensors are on the CPU and others are on the GPU, stacking fails. Move them to the same device first. The same rule applies to incompatible shapes.

It is also easy to insert the new dimension in the wrong position. If a model expects [batch, channels, height, width] and you stack with dim=1, the shape will be wrong even though the code runs.

Summary

  • Use torch.stack to turn a list of equal-shaped tensors into one tensor with dimension N + 1.
  • Use torch.cat only when you want to join tensors along an existing dimension.
  • Use unsqueeze for adding a singleton dimension to a single tensor.
  • Normalize variable-length inputs before stacking, often with padding.
  • Check shapes, devices, and dtypes before batching tensors.

Course illustration
Course illustration

All Rights Reserved.