def partition(X, group_sizes):
    """Partition the input according to the group_sizes at the last dimension"""
    List = []
    for i in range(X.size(-1) // group_sizes):
        List.append(X[..., i * group_sizes : (i + 1) * group_sizes])
    return List