code: def get_dataset(dataset_name):
      
      # mnist
      if dataset == "mnist":
          train_data = gluon.data.DataLoader(
              gluon.data.vision.MNIST('./data', train=True, transform=transformer),
              batch_size, shuffle=True, last_batch='discard')
  
          val_data = gluon.data.DataLoader(
              gluon.data.vision.MNIST('./data', train=False, transform=transformer),
              batch_size, shuffle=False)
      # cifar10
      elif dataset == "cifar10":
          train_data = gluon.data.DataLoader(
              gluon.data.vision.CIFAR10('./data', train=True, transform=transformer),
              batch_size, shuffle=True, last_batch='discard')
  
          val_data = gluon.data.DataLoader(
              gluon.data.vision.CIFAR10('./data', train=False, transform=transformer),
              batch_size, shuffle=False)
  
      return train_data, val_data
docs: Load the dataset and split it to train/valid data
  
      :param dataset_name: string
  
      Returns:
      train_data: int array
          training dataset
      val_data: int array
          valid dataset