def METHOD_NAME(x: torch.Tensor, split_size: Union[list, int]):
    r"""
    Transforms a tensor of shape (sum(Mi), K, L, ...) to N set of tensors of
    shape (Mi, K, L, ...) where Mi's are defined in split_size
    Args:
      x: tensor
      split_size: list, tuple or int defining the number of items for each tensor
        in the output list.
    Returns:
      x_list: A list of Tensors
    """
    return x.split(split_size, dim=0)