flatten() and unbind() in PyTorch

flatten() can remove zero or more dimensions from a 0D or more tensor as shown below:

import torch

my_tensor = torch.tensor(2)

torch.flatten(my_tensor)
torch.flatten(my_tensor, 0, 0)
torch.flatten(my_tensor, 0, -1)
torch.flatten(my_tensor, -1, 0)
torch.flatten(my_tensor, -1, -1)
# tensor([2])

my_tensor= torch.tensor([2, 7, 4]) # 1D tensor

torch.flatten(my_tensor)
torch.flatten(my_tensor, 0, 0)
torch.flatten(my_tensor, 0, -1)
torch.flatten(my_tensor, -1, 0)
torch.flatten(my_tensor, -1, -1)
# tensor([2, 7, 4])

my_tensor = torch.tensor([[2, 7, 4], [8, 3, 2]]) # 2D tensor

torch.flatten(my_tensor)
torch.flatten(my_tensor, 0, 0)
torch.flatten(my_tensor, 0, 1)
torch.flatten(my_tensor, 0, …

dimensions flat flatten import pytorch tensor torch