code: def load_data(traindir, valdir, args):
      # Data loading code
      print("Loading data")
      val_resize_size, val_crop_size, train_crop_size = (
          args.val_resize_size,
          args.val_crop_size,
          args.train_crop_size,
      )
      interpolation = InterpolationMode(args.interpolation)
  
      print("Loading training data")
      st = time.time()
      cache_path = _get_cache_path(traindir)
      if args.cache_dataset and os.path.exists(cache_path):
          # Attention, as the transforms are also cached!
          print(f"Loading dataset_train from {cache_path}")
          # TODO: this could probably be weights_only=True
          dataset, _ = torch.load(cache_path, weights_only=False)
      else:
          # We need a default value for the variables below because args may come
          # from train_quantization.py which doesn't define them.
          auto_augment_policy = getattr(args, "auto_augment", None)
          random_erase_prob = getattr(args, "random_erase", 0.0)
          ra_magnitude = getattr(args, "ra_magnitude", None)
          augmix_severity = getattr(args, "augmix_severity", None)
          dataset = torchvision.datasets.ImageFolder(
              traindir,
              presets.ClassificationPresetTrain(
                  crop_size=train_crop_size,
                  interpolation=interpolation,
                  auto_augment_policy=auto_augment_policy,
                  random_erase_prob=random_erase_prob,
                  ra_magnitude=ra_magnitude,
                  augmix_severity=augmix_severity,
                  backend=args.backend,
                  use_v2=args.use_v2,
              ),
          )
          if args.cache_dataset:
              print(f"Saving dataset_train to {cache_path}")
              utils.mkdir(os.path.dirname(cache_path))
              utils.save_on_master((dataset, traindir), cache_path)
      print("Took", time.time() - st)
  
      print("Loading validation data")
      cache_path = _get_cache_path(valdir)
      if args.cache_dataset and os.path.exists(cache_path):
          # Attention, as the transforms are also cached!
          print(f"Loading dataset_test from {cache_path}")
          # TODO: this could probably be weights_only=True
          dataset_test, _ = torch.load(cache_path, weights_only=False)
      else:
          if args.weights and args.test_only:
              weights = torchvision.models.get_weight(args.weights)
              preprocessing = weights.transforms(antialias=True)
              if args.backend == "tensor":
                  preprocessing = torchvision.transforms.Compose(
                      [
                          torchvision.transforms.PILToTensor(),
                          preprocessing,
                      ]
                  )
  
          else:
              preprocessing = presets.ClassificationPresetEval(
                  crop_size=val_crop_size,
                  resize_size=val_resize_size,
                  interpolation=interpolation,
                  backend=args.backend,
                  use_v2=args.use_v2,
              )
  
          dataset_test = torchvision.datasets.ImageFolder(
              valdir,
              preprocessing,
          )
          if args.cache_dataset:
              print(f"Saving dataset_test to {cache_path}")
              utils.mkdir(os.path.dirname(cache_path))
              utils.save_on_master((dataset_test, valdir), cache_path)
  
      print("Creating data loaders")
      if args.distributed:
          if hasattr(args, "ra_sampler") and args.ra_sampler:
              train_sampler = RASampler(
                  dataset, shuffle=True, repetitions=args.ra_reps
              )
          else:
              train_sampler = (
                  torch.utils.data.distributed.DistributedSampler(
                      dataset
                  )
              )
          test_sampler = (
              torch.utils.data.distributed.DistributedSampler(
                  dataset_test, shuffle=False
              )
          )
      else:
          train_sampler = torch.utils.data.RandomSampler(dataset)
          test_sampler = torch.utils.data.SequentialSampler(
              dataset_test
          )
  
      return dataset, dataset_test, train_sampler, test_sampler
  
imports: import datetime
  import os
  import time
  import warnings
  import presets
  import torch
  import torch.utils.data
  import torchvision
  import torchvision.transforms
  import utils
  from sampler import RASampler
  from torch import nn
  from torch.utils.data.dataloader import default_collate
  from torchvision.transforms.functional import InterpolationMode
  from transforms import get_mixup_cutmix
label_desc: Use image node_id 4 for calling a global function with example usage: load_data(traindir, valdir, args) and returns: dataset,  dataset_test,  train_sampler,  test_sampler