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)

EDIT: 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])

Want to know more cool tensor tricks? Check out my post on pairwise distance!