def __init__(
        self,
        data: AvalancheDataset,
        batch_size: int = 32,
        oversample_small_groups: bool = False,
        distributed_sampling: bool = True,
        **kwargs
    ):
        """Task-balanced data loader for Avalanche's datasets.

        The iterator returns a mini-batch balanced across each task, which
        makes it useful when training in multi-task scenarios whenever data is
        highly unbalanced.

        If `oversample_small_tasks == True` smaller tasks are
        oversampled to match the largest task. Otherwise, once the data for a
        specific task is terminated, that task will not be present in the
        subsequent mini-batches.

        :param data: an instance of `AvalancheDataset`.
        :param oversample_small_groups: whether smaller tasks should be
            oversampled to match the largest one.
        :param distributed_sampling: If True, apply the PyTorch
            :class:`DistributedSampler`. Defaults to True.
            Note: the distributed sampler is not applied if not running
            a distributed training, even when True is passed.
        :param kwargs: data loader arguments used to instantiate the loader for
            each task separately. See pytorch :class:`DataLoader`.
        """

        if "oversample_small_tasks" in kwargs:
            raise ValueError(
                "oversample_small_tasks is deprecated in favor of "
                "oversample_small_groups"
            )

        # Split data by task
        task_datasets = []
        task_labels_field = getattr(data, "targets_task_labels")
        assert isinstance(task_labels_field, DataAttribute)
        for task_label in task_labels_field.uniques:
            tidxs = task_labels_field.val_to_idx[task_label]
            tdata = data.subset(tidxs)
            task_datasets.append(tdata)

        super().__init__(
            task_datasets,
            oversample_small_groups=oversample_small_groups,
            batch_size=batch_size,
            distributed_sampling=distributed_sampling,
            **kwargs
        )
Task-balanced data loader for Avalanche's datasets.

        The iterator returns a mini-batch balanced across each task, which
        makes it useful when training in multi-task scenarios whenever data is
        highly unbalanced.

        If `oversample_small_tasks == True` smaller tasks are
        oversampled to match the largest task. Otherwise, once the data for a
        specific task is terminated, that task will not be present in the
        subsequent mini-batches.

        :param data: an instance of `AvalancheDataset`.
        :param oversample_small_groups: whether smaller tasks should be
            oversampled to match the largest one.
        :param distributed_sampling: If True, apply the PyTorch
            :class:`DistributedSampler`. Defaults to True.
            Note: the distributed sampler is not applied if not running
            a distributed training, even when True is passed.
        :param kwargs: data loader arguments used to instantiate the loader for
            each task separately. See pytorch :class:`DataLoader`.
        
avalanche/benchmarks/utils/data_loader.py