code: def prepare_dataloader(batch_size: int = 128):
      normalize = transforms.Normalize(
          (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
      )
      train_loader = DataLoader(
          datasets.CIFAR10(
              Path(__file__).parent / "data",
              train=True,
              transform=transforms.Compose(
                  [
                      transforms.RandomHorizontalFlip(),
                      transforms.RandomCrop(32, 4),
                      transforms.ToTensor(),
                      normalize,
                  ]
              ),
              download=True,
          ),
          batch_size=batch_size,
          shuffle=True,
          num_workers=8,
      )
  
      test_loader = DataLoader(
          datasets.CIFAR10(
              Path(__file__).parent / "data",
              train=False,
              transform=transforms.Compose(
                  [
                      transforms.ToTensor(),
                      normalize,
                  ]
              ),
          ),
          batch_size=batch_size,
          shuffle=False,
          num_workers=8,
      )
      return train_loader, test_loader
  
file_path: nni/examples/compression/models.py
imports: from pathlib import Path
  import torch
  import torch.nn.functional
  from torch.optim import Adam
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.data import DataLoader
  from torchvision import datasets, transforms
  from torchvision.models.mobilenetv3 import mobilenet_v3_small
  from torchvision.models.resnet import resnet18
  import nni
label_desc: Use image node_id 3 for calling a global function with example usage: prepare_dataloader(batch_size) and returns: train_loader,  test_loader