BucketIterator iterating more than the length

I have code like this:

batch_size = 100

train_dataset, test_dataset = convos.split(split_ratio=0.7)

train_iterator = torchtext.data.BucketIterator(
    train_dataset,
    batch_size=3000,
    sort_key=lambda x: torchtext.data.interleave_keys(len(x.context), len(x.response)),
    device=device
)

print("Batch size: ", batch_size)
print("Train size: ", len(train_iterator))

Batch size:  100
Train size:  1

Following this, I do

for count, batch in enumerate(train_iterator):
    print(count)

And this program doesn’t stop (at least I checked until count = 3000). I suppose it should iterate just once!

[mttk]

Yeah, I believe it is more or less agreed that the default infinite loop in Iterators is confusing.

I’d say that the expected behavior is to iterate for one epoch (so, the default would be repeat=False), and leave the infinite loop as an option if users prefer it. This might break something backwards but it’s manageable.

[zhangguanheng66]

I will just close this issue since a PR is attached. Thanks all for the help.