def test_profiler_closing(multiple_iterables):
    """Tests if the profiler terminates upon raising a StopIteration on an iterable dataset."""

    class TestDataset(IterableDataset):
        def __init__(self):
            self.list = list(range(1))

        def __iter__(self):
            return iter(self.list)

    fetcher = _PrefetchDataFetcher()
    if multiple_iterables:
        loader = CombinedLoader([DataLoader(TestDataset()), DataLoader(TestDataset())])
    else:
        loader = CombinedLoader(TestDataset())
    fetcher.setup(loader)
    profiler = SimpleProfiler()
    fetcher._start_profiler = lambda: profiler.start("test")
    fetcher._stop_profiler = lambda: profiler.stop("test")
    iter(fetcher)  # on epoch 0 start
    next(fetcher)  # raises StopIteration exception
    assert not bool(profiler.current_actions)
Tests if the profiler terminates upon raising a StopIteration on an iterable dataset.
tests/tests_pytorch/loops/test_fetchers.py