def test_inmem_batched_dataloader_shuffle_per_epoch(synthetic_dataset, reader_factory):
    """See if we are getting different shuffled result in different epochs"""
    epochs = 3
    with InMemBatchedDataLoader(reader_factory(synthetic_dataset.url, schema_fields=['^id$']),
                                batch_size=10, rows_capacity=100, num_epochs=epochs, shuffle=True) as loader:
        it = iter(loader)
        for i in range(epochs):
            # 10*10 == 100 rows in dataset
            cur_result = []
            for _ in range(10):
                batch = next(it)
                cur_result += list(batch['id'].numpy())
            if i == 0:
                first_result = cur_result
            elif i == 1:
                second_result = cur_result
                assert first_result != second_result
                assert set(first_result) == set(second_result)
            else:
                third_result = cur_result
                assert first_result != third_result
                assert second_result != third_result
                assert set(first_result) == set(third_result)
                assert set(second_result) == set(third_result)

        with pytest.raises(StopIteration):
            next(it)
See if we are getting different shuffled result in different epochs
petastorm/tests/test_pytorch_dataloader.py