code: def _imagenet_dataset(config):
      normalize = transforms.Normalize(
          mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
      )
      train_dir = os.path.join(config.data_dir, "train")
      test_dir = os.path.join(config.data_dir, "val")
      if hasattr(config, "use_aa") and config.use_aa:
          train_data = dset.ImageFolder(
              train_dir,
              transforms.Compose(
                  [
                      transforms.RandomResizedCrop(224),
                      transforms.RandomHorizontalFlip(),
                      ImageNetPolicy(),
                      transforms.ToTensor(),
                      normalize,
                  ]
              ),
          )
      else:
          train_data = dset.ImageFolder(
              train_dir,
              transforms.Compose(
                  [
                      transforms.RandomResizedCrop(224),
                      transforms.RandomHorizontalFlip(),
                      transforms.ColorJitter(
                          brightness=0.4,
                          contrast=0.4,
                          saturation=0.4,
                          hue=0.2,
                      ),
                      transforms.ToTensor(),
                      normalize,
                  ]
              ),
          )
  
      test_data = dset.ImageFolder(
          test_dir,
          transforms.Compose(
              [
                  transforms.Resize(256),
                  transforms.CenterCrop(224),
                  transforms.ToTensor(),
                  normalize,
              ]
          ),
      )
  
      return train_data, test_data
  
file_path: nni/examples/nas/legacy/cdarts/datasets/imagenet.py
imports: import os
  import numpy
  import torch
  import torchvision.datasets
  import torchvision.transforms
  from datasets.data_utils import ImageNetPolicy
  from datasets.data_utils import SubsetDistributedSampler
label_desc: Use image node_id 1 for calling a global function with example usage: _imagenet_dataset(config) and returns: train_data,  test_data