LeviViana:

torch.unbind


I think now torch.diagonal serves for same purpose.
import torch
a = np.array([[[1, 2, 3, 4],  # Input shape: (2, 3, 4)
               [5, 6, 7, 8],
               [9, 8, 7, 6]],
              [[5, 4, 3, 2],
               [1, 2, 3, 4],
               [5, 6, 7, 8]]])

a = torch.from_numpy(a)
# a.size() 
torch.diagonal(a, offset=0, dim1=-2, dim2=-1)
# torch.diagonal(b_tensor, offset=0, dim1=1, dim2=2)  # this return same results, set dim

return:
tensor([[1, 6, 7],
        [5, 2, 7]])

More details can be found in: https://pytorch.org/docs/master/torch.html?highlight=diag#torch.diagonal 26