def __next__(self):
    try:
        batch = next(self.dataset_iterator)
    except StopIteration:
        self.dataset_iterator = super().__iter__()
        batch = next(self.dataset_iterator)
    return batch
def __next__(self):
    try:
        batch = next(self.dataset_iterator)
    except StopIteration:
        self.dataset_iterator = super().__iter__()
        batch = next(self.dataset_iterator)
    return batch