def num_samples(self, dataset_split):
    """Determine the dataset sized given a dataset_split.

    Args:
      dataset_split: A problem.DatasetSplit.

    Returns:
      The desired number of samples for this dataset_split.
    """
    return {
        problem.DatasetSplit.TRAIN: 1000000,
        problem.DatasetSplit.EVAL: 10000,
        problem.DatasetSplit.TEST: 10000
    }[dataset_split]