code: def test_dataload_batch_balancing(self):
      benchmark = get_fast_benchmark()
      batch_size = 32
      replayPlugin = ReplayPlugin(mem_size=20)
  
      model = SimpleMLP(input_size=6, hidden_size=10)
      cl_strategy = Naive(
          model,
          SGD(
              model.parameters(),
              lr=0.001,
              momentum=0.9,
              weight_decay=0.001,
          ),
          CrossEntropyLoss(),
          train_mb_size=batch_size,
          train_epochs=1,
          eval_mb_size=100,
          plugins=[replayPlugin],
      )
      for step in benchmark.train_stream:
          adapted_dataset = step.dataset
          if len(replayPlugin.storage_policy.buffer) > 0:
              dataloader = ReplayDataLoader(
                  adapted_dataset,
                  replayPlugin.storage_policy.buffer,
                  oversample_small_tasks=True,
                  num_workers=0,
                  batch_size=batch_size,
                  shuffle=True,
              )
          else:
              dataloader = TaskBalancedDataLoader(adapted_dataset)
  
          for mini_batch in dataloader:
              mb_task_labels = mini_batch[-1]
              lengths = []
              for (
                  task_id
              ) in adapted_dataset.targets_task_labels.val_to_idx:
                  len_task = (mb_task_labels == task_id).sum()
                  lengths.append(len_task)
              if sum(lengths) == batch_size:
                  difference = max(lengths) - min(lengths)
                  self.assertLessEqual(difference, 1)
              self.assertLessEqual(sum(lengths), batch_size)
          cl_strategy.train(step)
  
imports: import unittest
  import torch
  from sklearn.datasets import make_classification
  from sklearn.model_selection import train_test_split
  from torch.optim import SGD
  from torch.nn import CrossEntropyLoss
  from torch.utils.data import TensorDataset, DataLoader
  from avalanche.benchmarks.utils import _concat_taskaware_classification_datasets
  from avalanche.benchmarks.utils.utils import concat_datasets
  from avalanche.models import SimpleMLP
  from avalanche.training.plugins import ReplayPlugin
  from avalanche.training.supervised import Naive
  from avalanche.benchmarks import nc_benchmark
  from avalanche.benchmarks.utils.data_loader import ReplayDataLoader, TaskBalancedDataLoader, GroupBalancedDataLoader, GroupBalancedInfiniteDataLoader
label_desc: Use image node_id 3 for calling the DataLoaderTests obj's underlying member method code with example usage: obj.test_dataload_batch_balancing() without return types