def _flatten_dataset_list(
    datasets: List[Union[FlatData[T_co], IDataset[T_co]]]
) -> List[IDataset[T_co]]:
    """Flatten the dataset tree if possible."""
    # Concat -> Concat branch
    # Flattens by borrowing the list of concatenated datasets
    # from the original datasets.
    flattened_list: List[IDataset[T_co]] = []
    for dataset in datasets:
        if len(dataset) == 0:
            continue
        elif (
            isinstance(dataset, FlatData)
            and dataset._indices is None
            and dataset._can_flatten
        ):
            flattened_list.extend(dataset._datasets)
        else:
            flattened_list.append(dataset)

    # merge consecutive Subsets if compatible
    new_data_list: List[IDataset[T_co]] = []
    for dataset in flattened_list:
        last_dataset = new_data_list[-1] if len(new_data_list) > 0 else None
        if (
            isinstance(dataset, FlatData)
            and len(new_data_list) > 0
            and isinstance(last_dataset, FlatData)
        ):
            new_data_list.pop()
            merged_ds = _maybe_merge_subsets(last_dataset, dataset)
            new_data_list.extend(merged_ds)
        elif (
            (dataset is not None)
            and len(new_data_list) > 0
            and (last_dataset is not None)
            and last_dataset is dataset
        ):
            new_data_list.pop()
            # the same dataset is repeated, using indices to avoid repeating it
            idxs = LazyIndices(
                LazyRange(0, len(last_dataset)), LazyRange(0, len(last_dataset))
            )
            merged_ds = [FlatData([last_dataset], indices=idxs)]
            new_data_list.extend(merged_ds)
        else:
            new_data_list.append(dataset)
    return new_data_list
Flatten the dataset tree if possible.
avalanche/benchmarks/utils/flat_data.py