Class TestIterator

dali/test/python/test_fw_iterators.py:1125–1153

class TestIterator:
    def __init__(self, iters_per_epoch, batch_size, total_iter_num=-1):
        self.n = iters_per_epoch
        self.total_n = total_iter_num
        self.batch_size = batch_size

    def __iter__(self):
        self.i = 0
        return self

    def __next__(self):
        batch = []
        # setting -1 means that no total iteration limit is set
        if self.i < self.n and self.total_n != 0:
            batch = [np.arange(0, 10, dtype=np.uint8) for _ in range(self.batch_size)]
            self.i += 1
            self.total_n -= 1
            return batch
        else:
            self.i = 0
            raise StopIteration

    next = __next__

    @property
    def size(
        self,
    ):
        return self.n * self.batch_size


@nottest