Adding a dimension to a tensor in PyTorchΒΆ

Updated on 2017-04-23

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.

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)

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)
>>> print(x1.size())
torch.Size([1, 10, 10])


comments powered by Disqus