def test_IterationBasedBatchSampler():
from torch.utils.data.sampler import SequentialSampler, BatchSampler
sampler = SequentialSampler([i for i in range(10)])
batch_sampler = BatchSampler(sampler, batch_size=2, drop_last=True)
batch_sampler = IterationBasedBatchSampler(batch_sampler, 5)
# check __len__
assert len(batch_sampler) == 5
for i, index in enumerate(batch_sampler):
assert [i * 2, i * 2 + 1] == index
# check start iter
batch_sampler.start_iter = 2
assert len(batch_sampler) == 3
