Adding a Dimension to a Tensor in PyTorch

Adding a dimension to a tensor can be important when you’re building deep learning models. In numpy, you can do this by inserting None into the axis you want to add.

import numpy as np

x1 = np.zeros((10, 10))
x2 = x1[None, :, :]

>>> print(x2.shape)
(1, 10, 10)

Update 2017-04-23: Good news! As of version 0.1.10, PyTorch supports None-style indexing. You should probably use that. But if you prefer to do it the old-fashioned way, read on.

Fortunately, it’s easy enough in PyTorch. Just pass the axis index into the .unsqueeze() method.

import torch

x1 = torch.zeros(10, 10)
x2 = x1.unsqueeze(0)

>>> print(x2.size())
torch.Size([1, 10, 10])

You can also do it in place using the underscore version .unsqueeze_().

x1 = torch.zeros(10, 10)
x1.unsqueeze_(0)

>>> print(x1.size())
torch.Size([1, 10, 10])