def gen_batches(self, batch_size: int, data_type: str = "train",
                    shuffle: bool = None) -> Iterator[Tuple[tuple, tuple]]:
        """
        Generates batches and expected output to train neural networks.
        If there are not enough samples from any task, samples are padded with None
        Args:
            batch_size: number of samples in batch
            data_type: can be either 'train', 'test', or 'valid'
            shuffle: whether to shuffle dataset before batching
        Yields:
            A tuple of a batch of inputs and a batch of expected outputs.
            Inputs and outputs are tuples. Element of inputs or outputs is a tuple which
            elements are x values of merged tasks in the order tasks are present in
            `tasks` argument of `__init__` method.
        """

        max_task_data_len = self.max_task_data_len[data_type]
        log.info(f'Batch size {batch_size} with gradient accumulation steps {self.gradient_accumulation_steps}')
        log.info(f'Efficient batch size {batch_size // self.gradient_accumulation_steps}')
        batch_size = batch_size // self.gradient_accumulation_steps

        if data_type == "train":
            generators = [
                SingleTaskBatchGenerator(iter_, batch_size, data_type, shuffle)
                for iter_ in self.task_iterators.values()
            ]
            # probs only required while training
            probs = self._get_probs("train")
            for step in range(self.steps_per_epoch):
                if (self.steps_taken + 1) % self.gradient_accumulation_steps == 0 or self.task_id is None:
                    self.task_id = np.random.choice(self.n_tasks, p=probs)
                x = [[None for _ in range(batch_size)] for _ in range(self.n_tasks)]
                y = [[None for _ in range(batch_size)] for _ in range(self.n_tasks)]
                x[self.task_id], y[self.task_id] = generators[self.task_id].__next__()
                if not all([s is None for s in x[self.task_id]]):
                    batch_to_yield = self._transform_before_yielding(
                        x, y, batch_size)
                    yield batch_to_yield

            self.epochs_done += 1
            # one additional step is taken while logging training metrics
            self.steps_taken -= 1
        else:
            eval_batch_size = 1
            x = [[None for _ in range(eval_batch_size)] for _ in range(self.n_tasks)]
            y = [[None for _ in range(eval_batch_size)] for _ in range(self.n_tasks)]
            generators = [
                SingleTaskBatchGenerator(
                    iter_, batch_size=eval_batch_size, data_type=data_type, shuffle=shuffle)
                for iter_ in self.task_iterators.values()
            ]
            for step in range(max_task_data_len):
                for task_id in range(self.n_tasks):
                    x[task_id], y[task_id] = generators[task_id].__next__()

                batches = self._transform_before_yielding(x, y, eval_batch_size)
                yield batches

        Generates batches and expected output to train neural networks.
        If there are not enough samples from any task, samples are padded with None
        Args:
            batch_size: number of samples in batch
            data_type: can be either 'train', 'test', or 'valid'
            shuffle: whether to shuffle dataset before batching
        Yields:
            A tuple of a batch of inputs and a batch of expected outputs.
            Inputs and outputs are tuples. Element of inputs or outputs is a tuple which
            elements are x values of merged tasks in the order tasks are present in
            `tasks` argument of `__init__` method.
        
deeppavlov/dataset_iterators/multitask_iterator.py