def batch_split(batch_size, max_num):
    """Split into equal parts of {batch_size} as well as the tail"""

    if batch_size > max_num:
        print("Fix the batch size to maximum number.")
        batch_size = max_num

    max_range = list(range(max_num))
    num_splits = max_num // batch_size
    num_splits = (
        num_splits - 1 if max_num % batch_size == 0 else num_splits
    )  # for edge case, there will be an empty batch
    splits = []
    for i in range(num_splits + 1):
        start = i * batch_size
        end = min((i + 1) * batch_size, max_num)  # dealing the tail part
        splits.append(max_range[start:end])
    assert len(splits) == num_splits + 1
    return splits