code: def test_all_dataset_batches_restart(self) -> None:
      dataloader_1 = DataLoader(RandomDataset(32, 8), batch_size=8)
      dataloader_2 = DataLoader(RandomDataset(32, 16), batch_size=8)
  
      all_dataset_batches = AllDatasetBatches(
          StoppingMechanism.RESTART_UNTIL_ALL_DATASETS_EXHAUSTED
      )
  
      multi_dataloader = iter(
          MultiDataLoader(
              self._get_dataloaders_dict(dataloader_1, dataloader_2),
              all_dataset_batches,
          )
      )
  
      # batch should contain both "1" and "2" keys
      batch = next(multi_dataloader)
      self.assertTrue(len(batch) == 2)
      self.assertTrue(("1" in batch) and ("2" in batch))
  
      # Dataset "1" was finished. batch now contains only samples from "2"
      batch = next(multi_dataloader)
      self.assertTrue(len(batch) == 2)
      self.assertTrue(("1" in batch) and ("2" in batch))
  
      # StopIteration after both datasets have been exhausted
      with self.assertRaises(StopIteration):
          batch = next(multi_dataloader)
  
imports: import random
  import unittest
  from collections import Counter
  from typing import Any, cast, Dict, Iterable, Iterator, List, Mapping, Union
  import torch
  from torch.utils.data import DataLoader, Dataset
  from torchtnt.utils.data.iterators import AllDatasetBatches, DataIterationStrategy, InOrder, MultiIterator, RandomizedBatchSampler, RoundRobin, StoppingMechanism
  from torchtnt.utils.data.multi_dataloader import MultiDataLoader
label_desc: Use image node_id 4 for calling the TestMultiDataLoader obj's underlying member method code with example usage: obj.test_all_dataset_batches_restart() without return types