def split_validation_class_balanced(
    validation_size: Union[int, float],
    dataset: TaskAwareSupervisedClassificationDataset,
) -> Tuple[
    TaskAwareSupervisedClassificationDataset, TaskAwareSupervisedClassificationDataset
]:
    """Class-balanced dataset split.

    This splitting strategy splits `dataset` into train and validation data of
    size `validation_size` using a class-balanced split.
    Samples of each class are chosen randomly.

    You can use this split strategy to split a benchmark with::

        validation_size = 0.2
        foo = lambda data: split_validation_class_balanced(validation_size, data)
        bm = benchmark_with_validation_stream(bm, split_strategy=foo)

    :param validation_size: The percentage of samples to allocate to the
        validation experience as a float between 0 and 1.
    :param dataset: The dataset to split.
    :return: A tuple containing 2 elements: the new training and validation
        datasets.
    """
    if not isinstance(validation_size, float):
        raise ValueError("validation_size must be an integer")
    if not 0.0 <= validation_size <= 1.0:
        raise ValueError("validation_size must be a float in [0, 1].")

    if validation_size > len(dataset):
        raise ValueError(
            f"Can't create the validation experience: not enough "
            f"instances. Required {validation_size}, got only"
            f"{len(dataset)}"
        )
    exp_indices = list(range(len(dataset)))
    targets_as_tensor = torch.as_tensor(dataset.targets)
    exp_classes: List[int] = targets_as_tensor.unique().tolist()

    # shuffle exp_indices
    exp_indices_t = torch.as_tensor(exp_indices)[torch.randperm(len(exp_indices))]
    # shuffle the targets as well
    exp_targets = targets_as_tensor[exp_indices_t]

    train_exp_indices: list[int] = []
    valid_exp_indices: list[int] = []
    for cid in exp_classes:  # split indices for each class separately.
        c_indices = exp_indices_t[exp_targets == cid]
        valid_n_instances = int(validation_size * len(c_indices))
        valid_exp_indices.extend(c_indices[:valid_n_instances])
        train_exp_indices.extend(c_indices[valid_n_instances:])

    result_train_dataset = dataset.subset(train_exp_indices)
    result_valid_dataset = dataset.subset(valid_exp_indices)
    return result_train_dataset, result_valid_dataset
Class-balanced dataset split.

    This splitting strategy splits `dataset` into train and validation data of
    size `validation_size` using a class-balanced split.
    Samples of each class are chosen randomly.

    You can use this split strategy to split a benchmark with::

        validation_size = 0.2
        foo = lambda data: split_validation_class_balanced(validation_size, data)
        bm = benchmark_with_validation_stream(bm, split_strategy=foo)

    :param validation_size: The percentage of samples to allocate to the
        validation experience as a float between 0 and 1.
    :param dataset: The dataset to split.
    :return: A tuple containing 2 elements: the new training and validation
        datasets.
    
avalanche/benchmarks/scenarios/dataset_scenario.py