class_docstr: Custom data loader for rehearsal/replay strategies.
class_docstr_tok: ["Custom","data","loader","for","rehearsal\/replay","strategies","."]
code: def _get_datasets_and_batch_sizes(
      data: AvalancheDataset,
      batch_sizes_def: Union[int, Dict[int, int]],
      remaining_examples: int,
      task_balanced_dataloader: bool,
  ):
      datasets: List[AvalancheDataset] = []
      batch_sizes: List[int] = []
      batch_size_per_task = not isinstance(batch_sizes_def, int)
  
      if task_balanced_dataloader or batch_size_per_task:
          for task_id in data.task_set:
              dataset = data.task_set[task_id]
  
              if batch_size_per_task:
                  current_batch_size = batch_sizes_def[task_id]
              else:
                  current_batch_size = batch_sizes_def
  
              if remaining_examples > 0:
                  current_batch_size += 1
                  remaining_examples -= 1
  
              datasets.append(dataset)
              batch_sizes.append(current_batch_size)
      else:
          # Current data is loaded without task balancing
          datasets.append(data)
          batch_sizes.append(batch_sizes_def)
      return batch_sizes, datasets
  
imports: from typing import Any, Dict, Iterator, List, Optional, Sequence, Sized, Union
  import numpy
  from torch.utils.data import DistributedSampler, Dataset
  from torch.utils.data.dataloader import DataLoader
  from avalanche.benchmarks.utils.data import AvalancheDataset
  from avalanche.benchmarks.utils.data_attribute import DataAttribute
  from avalanche.benchmarks.utils.ffcv_support.ffcv_components import HybridFfcvLoader, has_ffcv_support
  from avalanche.distributed.distributed_helper import DistributedHelper
  from torch.utils.data.sampler import Sampler, BatchSampler
  from torch.utils.data import ConcatDataset
label_desc: Use image node_id 2 for calling the ReplayDataLoader obj's underlying member method code with example usage: obj._get_datasets_and_batch_sizes(data, batch_sizes_def, remaining_examples, task_balanced_dataloader) and returns: batch_sizes,  datasets