


x = torch.arange(14).reshape(2, 7)
      >>> x
      tensor([[ 0,  1,  2,  3,  4,  5,  6],
              [ 7,  8,  9, 10, 11, 12, 13]])
      >>> torch.tensor_split(x, 3, dim=1)
      (tensor([[0, 1, 2],
              [7, 8, 9]]),
       tensor([[ 3,  4],
              [10, 11]]),
       tensor([[ 5,  6],
              [12, 13]]))
      >>> torch.tensor_split(x, (1, 6), dim=1)
      (tensor([[0],
              [7]]),
       tensor([[ 1,  2,  3,  4,  5],
              [ 8,  9, 10, 11, 12]]),
       tensor([[ 6],
              [13]]))


