Function check_stop_iter

dali/test/python/test_fw_iterators.py:1167–1190

(
    fw_iter, iterator_name, batch_size, epochs, iter_num, total_iter_num, auto_reset, infinite
)

def check_stop_iter(
    fw_iter, iterator_name, batch_size, epochs, iter_num, total_iter_num, auto_reset, infinite
):
    it = TestIterator(iter_num, batch_size, total_iter_num)
    pipe = create_test_iter_pipeline(batch_size, 0, it)
    if infinite:
        iter_size = -1
    else:
        iter_size = it.size
    loader = fw_iter(pipe, iter_size, auto_reset)
    count = 0
    for _ in range(epochs):
        for _ in enumerate(loader):
            count += 1
        if not auto_reset:
            loader.reset()
    if total_iter_num < 0:
        # infinite source of data
        assert count == iter_num * epochs
    else:
        # at most total_iter_num should be returned by the iterator
        assert count == min(
            total_iter_num, iter_num * epochs
        ), f"{count}, {total_iter_num}, {iter_num} * {epochs} == {iter_num * epochs}"


@raises(Exception, glob="Negative size is supported only for a single pipeline")