def __init__(self, source_iterator: CheckpointableIterator, batch_size: int):
        """
        Args:
            source_iterator: checkpointable input iterators
            batch_size: number of items per batch
        """
        if not isinstance(source_iterator, CheckpointableIterator):
            raise ValueError('source_iterator has to be a CheckpointableIterator')
        if batch_size <= 0:
            raise ValueError('batch_size has to be positive')
        self._source_iterator = source_iterator  # type: CheckpointableIterator
        self._batch_size = batch_size            # type: int
        self.setstate(None)

        Args:
            source_iterator: checkpointable input iterators
            batch_size: number of items per batch
        
kosmos-2/infinibatch/infinibatch/iterators.py