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
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
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!