    def count_batches(self, dataset_counts: Dict[str, int]) -> int:
        """
        Given the number of instances per dataset, this returns the total number of batches
        the scheduler will return.
        """
        raise NotImplementedError
Given the number of instances per dataset, this returns the total number of batches
the scheduler will return.
allennlp/data/data_loaders/multitask_scheduler.py