Adding a Dimension to a Tensor in PyTorch

Posted 2017-03-09 • Updated 2020-01-02

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!

Leave a Comment

Your email address will not be published. Required fields are marked *