if self.batched:
            if self._state_dict:
                self._state_dict["previous_state"] = self.ex_iterable.state_dict()
                self._state_dict["num_examples_since_previous_state"] = 0
                self._state_dict["previous_state_example_idx"] = current_idx
            for key, example in iterator:
                # If `batched`, first build the batch, if `batch_size` is None or <=0, then the batch is the whole dataset
                iterator_batch = (
                    iterator
                    if self.batch_size is None or self.batch_size <= 0
                    else islice(iterator, self.batch_size - 1)
                )
                key_examples_list = [(key, example)] + list(iterator_batch)
                keys, examples = zip(*key_examples_list)
                batch = _examples_to_batch(examples)
                batch = format_dict(batch) if format_dict else batch
                # then compute the mask for the batch
                inputs = batch