code: def get_cifar_vit(
      dataset_name,
      data_dir,
      split,
      transform=None,
      imsize=None,
      bucket="pytorch-data",
      **kwargs,
  ):
      if imsize == 224:
          if split == "train":
              transform_data = transforms.Compose(
                  [  # transforms.ColorJitter(brightness= 0.4, contrast= 0.4, saturation= 0.4, hue= 0.1),
                      transforms.Resize(256),
                      transforms.RandomCrop(224),
                      transforms.RandomHorizontalFlip(),
                      transforms.ToTensor(),
                      transforms.Normalize(
                          (0.5071, 0.4867, 0.4408),
                          (0.2675, 0.2565, 0.2761),
                      ),
                  ]
              )
          else:
              transform_data = transforms.Compose(
                  [
                      transforms.Resize(256),
                      transforms.CenterCrop(224),
                      transforms.ToTensor(),
                      transforms.Normalize(
                          (0.5071, 0.4867, 0.4408),
                          (0.2675, 0.2565, 0.2761),
                      ),
                  ]
              )
          if dataset_name == "cifar10":
              return datasets.CIFAR10(
                  data_dir,
                  train=(split == "train"),
                  transform=transform_data,
                  download=True,
                  **kwargs,
              )
          elif dataset_name == "cifar100":
              return datasets.CIFAR100(
                  data_dir,
                  train=(split == "train"),
                  transform=transform_data,
                  download=True,
                  **kwargs,
              )
          else:
              assert dataset_name in ["cifar10", "cifar100"]
      else:
          if split == "train":
              transform_data = transforms.Compose(
                  [  # transforms.ColorJitter(brightness= 0.4, contrast= 0.4, saturation= 0.4, hue= 0.1),
                      transforms.Resize(imsize),
                      transforms.RandomCrop(
                          imsize, padding=(imsize // 8)
                      ),
                      transforms.RandomHorizontalFlip(),
                      transforms.ToTensor(),
                      transforms.Normalize(
                          (0.5071, 0.4867, 0.4408),
                          (0.2675, 0.2565, 0.2761),
                      ),
                  ]
              )
          else:
              transform_data = transforms.Compose(
                  [
                      transforms.Resize(imsize),
                      transforms.ToTensor(),
                      transforms.Normalize(
                          (0.5071, 0.4867, 0.4408),
                          (0.2675, 0.2565, 0.2761),
                      ),
                  ]
              )
          if dataset_name == "cifar10":
              return datasets.CIFAR10(
                  data_dir,
                  train=(split == "train"),
                  transform=transform_data,
                  download=True,
                  **kwargs,
              )
          elif dataset_name == "cifar100":
              return datasets.CIFAR100(
                  data_dir,
                  train=(split == "train"),
                  transform=transform_data,
                  download=True,
                  **kwargs,
              )
          else:
              assert dataset_name in ["cifar10", "cifar100"]
  
file_path: deepspeedexamples/training/data_efficiency/vit_finetuning/utils/get_data.py
imports: import torch
  import os
  import torchvision.transforms
  import torchvision.datasets
label_desc: Use image node_id 7 for calling a global function with example usage: get_cifar_vit(dataset_name, data_dir, split, transform, imsize, bucket) and returns: datasets, datasets, datasets, datasets