code: def test_distil_label_patcher(uid_dataset_cls: str, storage: str):
      dataset = DemoDataset()
  
      if uid_dataset_cls == "index":
          uid_dataset = IndexedDataset(dataset=dataset)
      elif uid_dataset_cls == "hash":
          uid_dataset = HashedDataset(dataset=dataset, hash_fn=hash)
      elif uid_dataset_cls == "augumentation":
          uid_dataset = AugmentationDataset(
              IndexedDataset(dataset=dataset), transform=transform_fn
          )
  
      if storage == "memory":
          patcher = DistilLabelPatcher(
              teacher_predict=teacher_predict,
              dataset=uid_dataset,
              create_dataloader=create_dataloader,
              keep_in_memory=True,
          )
      elif storage == "file":
          patcher = DistilLabelPatcher(
              teacher_predict=teacher_predict,
              dataset=uid_dataset,
              create_dataloader=create_dataloader,
              keep_in_memory=False,
              cache_folder=log_dir,
              cache_mode="pickle",
          )
      # add this storage back after it is done
      # elif storage == 'hdf5':
      #     patcher = DistilLabelPatcher(teacher_predict=teacher_predict, dataset=uid_dataset, create_dataloader=create_dataloader,
      #                                  keep_in_memory=False, cache_folder=log_dir, cache_mode='hdf5')
  
      patcher.generate_distillation_labels()
      distil_dataloader = patcher.create_patched_dataloader()
  
      if not uid_dataset_cls == "augumentation":
          dataloader = create_dataloader(dataset)
          for d1, d2 in zip(dataloader, distil_dataloader):
              assert torch.equal(teacher_predict(d1), d2[0])
      else:
          dataloader = create_augmentation_dataloader(uid_dataset)
          for d1, d2 in zip(dataloader, distil_dataloader):
              assert torch.equal(teacher_predict(d1), d2[0])
  
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 6 for calling a global function with example usage: test_distil_label_patcher(uid_dataset_cls, storage) without return types