def test_with_torch_api(synthetic_dataset):
    """Verify that WeightedSamplingReader is compatible with petastorm.pytorch.DataLoader"""
    readers = [reader0, reader1]

    with WeightedSamplingReader(readers, [0.5, 0.5]) as mixer:
        assert not mixer.batched_output
        sample = next(mixer)
        assert sample is not None
        with DataLoader(mixer, batch_size=2) as loader:
            for batch in loader:
                assert batch['f1'].shape[0] == 2
                break
Verify that WeightedSamplingReader is compatible with petastorm.pytorch.DataLoader
petastorm/tests/test_weighted_sampling_reader.py