code: def load_dataset(name, root, sample="default", **kwargs):
      """
      Default dataset wrapper
  
      :param name (string): Name of the dataset (Out of cifar10/100, imagenet, tinyimagenet, CUB200, STANFORD120, MIT67).
      :param root (string): Path to download the dataset.
      :param sample (string): Default (random) sampling as the classic pytorch dataloader or Pairwise sampling as mentioned in
      the paper "Regularizing Class-wise Predictions via Self-knowledge Distillation"
      """
  
      # Dataset
      if name in [
          "imagenet",
          "tinyimagenet",
          "CUB200",
          "STANFORD120",
          "MIT67",
      ]:
          if name == "tinyimagenet":
              transform_train = transforms.Compose(
                  [
                      transforms.RandomResizedCrop(32),
                      transforms.RandomHorizontalFlip(),
                      transforms.ToTensor(),
                      transforms.Normalize(
                          (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
                      ),
                  ]
              )
              transform_test = transforms.Compose(
                  [
                      transforms.Resize(32),
                      transforms.ToTensor(),
                      transforms.Normalize(
                          (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
                      ),
                  ]
              )
  
              train_val_dataset_dir = os.path.join(root, "train")
              test_dataset_dir = os.path.join(root, "val")
  
              trainset = DatasetWrapper(
                  datasets.ImageFolder(
                      root=train_val_dataset_dir,
                      transform=transform_train,
                  )
              )
              valset = DatasetWrapper(
                  datasets.ImageFolder(
                      root=test_dataset_dir, transform=transform_test
                  )
              )
  
          elif name == "imagenet":
              transform_train = transforms.Compose(
                  [
                      transforms.RandomResizedCrop(224),
                      transforms.RandomHorizontalFlip(),
                      transforms.ToTensor(),
                      transforms.Normalize(
                          (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
                      ),
                  ]
              )
              transform_test = transforms.Compose(
                  [
                      transforms.Resize(256),
                      transforms.CenterCrop(224),
                      transforms.ToTensor(),
                      transforms.Normalize(
                          (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
                      ),
                  ]
              )
              train_val_dataset_dir = os.path.join(root, "train")
              test_dataset_dir = os.path.join(root, "val")
  
              trainset = DatasetWrapper(
                  datasets.ImageFolder(
                      root=train_val_dataset_dir,
                      transform=transform_train,
                  )
              )
              valset = DatasetWrapper(
                  datasets.ImageFolder(
                      root=test_dataset_dir, transform=transform_test
                  )
              )
  
          else:
              transform_train = transforms.Compose(
                  [
                      transforms.RandomResizedCrop(224),
                      transforms.RandomHorizontalFlip(),
                      transforms.ToTensor(),
                      transforms.Normalize(
                          (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
                      ),
                  ]
              )
              transform_test = transforms.Compose(
                  [
                      transforms.Resize(256),
                      transforms.CenterCrop(224),
                      transforms.ToTensor(),
                      transforms.Normalize(
                          (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
                      ),
                  ]
              )
  
              train_val_dataset_dir = os.path.join(root, name, "train")
              test_dataset_dir = os.path.join(root, name, "test")
  
              trainset = DatasetWrapper(
                  datasets.ImageFolder(
                      root=train_val_dataset_dir,
                      transform=transform_train,
                  )
              )
              valset = DatasetWrapper(
                  datasets.ImageFolder(
                      root=test_dataset_dir, transform=transform_test
                  )
              )
  
      elif name.startswith("cifar"):
          transform_train = transforms.Compose(
              [
                  transforms.RandomCrop(32, padding=4),
                  transforms.RandomHorizontalFlip(),
                  transforms.ToTensor(),
                  transforms.Normalize(
                      (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                  ),
              ]
          )
          transform_test = transforms.Compose(
              [
                  transforms.ToTensor(),
                  transforms.Normalize(
                      (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                  ),
              ]
          )
  
          if name == "cifar10":
              CIFAR = datasets.CIFAR10
          else:
              CIFAR = datasets.CIFAR100
  
          trainset = DatasetWrapper(
              CIFAR(
                  root,
                  train=True,
                  download=True,
                  transform=transform_train,
              )
          )
          valset = DatasetWrapper(
              CIFAR(
                  root,
                  train=False,
                  download=True,
                  transform=transform_test,
              )
          )
      else:
          raise Exception("Unknown dataset: {}".format(name))
  
      # Sampler
      if sample == "default":
          get_train_sampler = lambda d: BatchSampler(
              RandomSampler(d), kwargs["batch_size"], False
          )
          get_test_sampler = lambda d: BatchSampler(
              SequentialSampler(d), kwargs["batch_size"], False
          )
  
      elif sample == "pair":
          get_train_sampler = lambda d: PairBatchSampler(
              d, kwargs["batch_size"]
          )
          get_test_sampler = lambda d: BatchSampler(
              SequentialSampler(d), kwargs["batch_size"], False
          )
  
      else:
          raise Exception("Unknown sampling: {}".format(sampling))
  
      train_loader = DataLoader(
          trainset,
          batch_sampler=get_train_sampler(trainset),
          num_workers=4,
      )
      val_loader = DataLoader(
          valset, batch_sampler=get_test_sampler(valset), num_workers=4
      )
  
      return train_loader, val_loader
  
imports: import csv
  import os
  import random
  from collections import defaultdict
  import numpy
  import torchvision
  from PIL import Image
  from torch.utils.data import BatchSampler, DataLoader, Dataset, RandomSampler, Sampler, SequentialSampler, Subset
  from torchvision import datasets, transforms
label_desc: Use image node_id 1 for calling a global function with example usage: load_dataset(name, root, sample) and returns: train_loader,  val_loader