text: """
  dataset and data loader
  """
  import os
  import torch
  import torchvision.transforms as transforms
  from torch.utils.data import DataLoader
  from torchvision import datasets
  
  classes = {"cubs_cropped":200, "flowers":102, "sketches":250, "stanford_cars_cropped":196, "wikiart":195}
  
  def get_loader(args):
      # Preparing data
      if args.dataset == 'cifar10':
          data_path = args.data_path + 'cifar'
  
          mean = [x / 255 for x in [125.3, 123.0, 113.9]]
          std = [x / 255 for x in [63.0, 62.1, 66.7]]
  
          train_transform = transforms.Compose([
              transforms.RandomHorizontalFlip(),
              transforms.RandomCrop(32, padding=4),
              transforms.ToTensor(),
              transforms.Normalize(mean, std)
          ])
          test_transform = transforms.Compose([
              transforms.ToTensor(),
              transforms.Normalize(mean, std)
          ])
  
          trainset = datasets.CIFAR10(root=args.data_path, train=True, download=True, transform=train_transform)
          trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=2)
  
          testset = datasets.CIFAR10(root=args.data_path, train=False, download=True, transform=test_transform)
          testloader = torch.utils.data.DataLoader(testset, batch_size=args.batch_size, shuffle=False, num_workers=2)
          num_classes = 10
      elif args.dataset == 'imagenet':
          mean = [0.485, 0.456, 0.406]
          std = [0.229, 0.224, 0.225]
  
          train_transform = transforms.Compose([
              transforms.RandomResizedCrop(224),
              transforms.RandomHorizontalFlip(),
              transforms.ToTensor(),
              transforms.Normalize(mean, std)
          ])
          test_transform = transforms.Compose([
              transforms.Scale(256),
              transforms.CenterCrop(224),
              transforms.ToTensor(),
              transforms.Normalize(mean, std)
          ])  # here is actually the validation dataset
  
          train_dir = os.path.join(args.data_path, 'train')
          test_dir = os.path.join(args.data_path, 'val')
  
          train_data = datasets.ImageFolder(train_dir, transform=train_transform)
          test_data = datasets.ImageFolder(test_dir, transform=test_transform)
  
          trainloader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True)
          testloader = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True)
          num_classes = 1000
      elif args.dataset in classes.keys():
          mean=[0.485, 0.456, 0.406]
          std=[0.229, 0.224, 0.225]
  
          train_transform = transforms.Compose([
              transforms.Scale(256),
              transforms.RandomSizedCrop(224),
              transforms.RandomHorizontalFlip(),
              transforms.ToTensor(),
              transforms.Normalize(mean, std)
          ])
  
          test_transform = transforms.Compose([
              transforms.Scale(256),
              transforms.CenterCrop(224),
              transforms.ToTensor(),
              transforms.Normalize(mean, std)
          ])
          
          train_dir = os.path.join(args.data_path, 'train')
          test_dir = os.path.join(args.data_path, 'test')
  
          train_data = datasets.ImageFolder(train_dir, transform=train_transform)
          test_data = datasets.ImageFolder(test_dir, transform=test_transform)
  
          trainloader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True)
          testloader = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True)
          num_classes = classes[args.dataset]
      else:
          raise ValueError("Unrecegonized dataset!")
      
      return trainloader, testloader, num_classes
  