pytorch unsqueeze code example

Example 1: pytorch tensor add one dimension

# ADD ONE DIMENSION: .unsqueeze(dim)

my_tensor = torch.tensor([1,3,4])
# tensor([1,3,4])

my_tensor.unsqueeze(0)
# tensor([[1,3,4]])

my_tensor.unsqueeze(1)
# tensor([[1],
#         [3],
#         [4]])

Example 2: pytorch squeeze

x = torch.zeros(2, 1, 2, 1, 2)
x.size()
>>> torch.Size([2, 1, 2, 1, 2])

y = torch.squeeze(x) # remove 1
y.size()
>>> torch.Size([2, 2, 2])

y = torch.squeeze(x, 0)
y.size()
>>> torch.Size([2, 1, 2, 1, 2])

y = torch.squeeze(x, 1)
y.size()
>>> torch.Size([2, 2, 1, 2])

Example 3: pytorch unsqueeze

ft = torch.Tensor([0, 1, 2])
print(ft.shape)
>>> torch.Size([3])

print(ft.unsqueeze(0)) # 0 means first dimension
print(ft.unsqueeze(0).shape)
>>> tensor([[0., 1., 2.]])
>>> torch.Size([1, 3])