code: def test_balanced_small_batches():
      sampler = BalancedBatchSampler(MyDataset(), batch_size=4 * 4)
      assert len(sampler) == 3
      batches = list(itertools.islice(sampler, 8))
      assert len(batches) == 3
      for batch in batches:
          assert len(batch) == 16
          counts = Counter(map(idx_to_class, batch))
          assert counts == {"1": 4, "2": 4, "3": 4, "4": 4}
  
      # no duplicates in the 2 first batches
      first_batches = batches[0] + batches[1]
      assert len(first_batches) == len(set(first_batches))
  
      # duplicates for classes 1, 2, 4 in the last batch
      for i in range(2, 3):
          first_batches += batches[i]
      counts = Counter(first_batches)
      assert {
          idx_to_class(idx) for idx, nb in counts.items() if nb > 1
      } == {"1", "2", "4"}
  
file_path: pykale/tests/loaddata/test_sampler.py
imports: import itertools
  from collections import Counter
  import numpy
  from kale.loaddata.sampler import BalancedBatchSampler, ReweightedBatchSampler
label_desc: Use image node_id 2 for calling a global function with example usage: test_balanced_small_batches() without return types