split() and vsplit() in PyTorch

split() can split a 1D or more D tensor into 1 or more tensors as shown below. *Setting a dimension to the 2nd argument can select the split position of a tensor:

import torch

my_tensor = torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]])

torch.split(my_tensor, 1)
my_tensor.split(1)
# (tensor([[0, 1, 2, 3]]),
#  tensor([[4, 5, 6, 7]]),
#  tensor([[8, 9, 10, 11]]))

torch.split(my_tensor, 2)
my_tensor.split(2)
# (tensor([[0, 1, 2, 3], [4, 5, 6, 7]]),
#  tensor([[8, …

functional import pytorch split tensor torch