code: def create_augmentation_dataloader(dataset: AugmentationDataset):
      suid_dataset = dataset._dataset
  
      def collate_fn(data):
          batch = []
          cpu_rng_state = torch.random.get_rng_state()
          for uid, sample in data:
              seed = dataset._suid_seed[str(uid)][-1]
              torch.random.manual_seed(seed)
              batch.append(transform_fn(sample))
          torch.random.set_rng_state(cpu_rng_state)
          return torch.stack(batch)
  
      return DataLoader(
          dataset=suid_dataset, batch_size=10, collate_fn=collate_fn
      )
  
file_path: nni/test/algo/compression/distillation/test_distil_label_patcher.py
imports: from __future__ import annotations
  import hashlib
  import pytest
  import torch
  from torch.utils.data import Dataset, DataLoader
  from nni.contrib.distillation.label_patcher import DistilLabelPatcher
  from nni.contrib.distillation.uid_dataset import IndexedDataset, HashedDataset, AugmentationDataset
  from ..assets.common import log_dir
label_desc: Use image node_id 3 for calling a global function with example usage: create_augmentation_dataloader(dataset) and returns: DataLoader, torch