code: def get_pytorch_train_loader(
      data_path,
      image_size,
      batch_size,
      num_classes,
      one_hot,
      interpolation="bilinear",
      augmentation=None,
      start_epoch=0,
      workers=5,
      _worker_init_fn=None,
      prefetch_factor=2,
      memory_format=torch.contiguous_format,
  ):
      interpolation = {
          "bicubic": Image.BICUBIC,
          "bilinear": Image.BILINEAR,
      }[interpolation]
      traindir = os.path.join(data_path, "train")
      transforms_list = [
          transforms.RandomResizedCrop(
              image_size, interpolation=interpolation
          ),
          transforms.RandomHorizontalFlip(),
      ]
      if augmentation == "disabled":
          pass
      elif augmentation == "autoaugment":
          transforms_list.append(AutoaugmentImageNetPolicy())
      else:
          raise NotImplementedError(
              f"Automatic augmentation: '{augmentation}' is not supported"
              " for PyTorch data loader."
          )
      train_dataset = datasets.ImageFolder(
          traindir, transforms.Compose(transforms_list)
      )
  
      if torch.distributed.is_initialized():
          train_sampler = (
              torch.utils.data.distributed.DistributedSampler(
                  train_dataset, shuffle=True
              )
          )
      else:
          train_sampler = None
  
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          sampler=train_sampler,
          batch_size=batch_size,
          shuffle=(train_sampler is None),
          num_workers=workers,
          worker_init_fn=_worker_init_fn,
          pin_memory=True,
          collate_fn=partial(fast_collate, memory_format),
          drop_last=True,
          persistent_workers=True,
          prefetch_factor=prefetch_factor,
      )
  
      return (
          PrefetchedWrapper(
              train_loader, start_epoch, num_classes, one_hot
          ),
          len(train_loader),
      )
  
imports: import os
  import torch
  import numpy
  from PIL import Image
  from functools import partial
  from image_classification.autoaugment import AutoaugmentImageNetPolicy
  import torchvision.datasets
  import torchvision.transforms
label_desc: Use image node_id 6 for calling a global function with example usage: get_pytorch_train_loader(data_path, image_size, batch_size, num_classes, one_hot, interpolation, augmentation, start_epoch, workers, _worker_init_fn, prefetch_factor, memory_format) and returns: PrefetchedWrapper,  len