code: def num_samples(self, dataset_split):
      
      return {
          problem.DatasetSplit.TRAIN: 1000000,
          problem.DatasetSplit.EVAL: 10000,
          problem.DatasetSplit.TEST: 10000
      }[dataset_split]
docs: Determine the dataset sized given a dataset_split.
  
      Args:
        dataset_split: A problem.DatasetSplit.
  
      Returns:
        The desired number of samples for this dataset_split.