Length error using `map` with datasets

sachin
Mon, 29 May 2023 08:01:21 +0000

Hey, almost there. Since laion_ds was an IterableDataset class I could only do the following: laion_ds_batched = laion_ds.map(collate_fn, batched=True, batch_size=8, remove_columns=next(iter(laion_ds)).keys()). However, when I do x = next(iter(laion_ds_batched)) the returned dictionary only returns one element out the possible 8. The stranger thing is when I set a breakpoint within the CollateFn class, it is batching more than one element properly.

A minimal working example can be found here. Is this possibly a bug? laion_hf_dataset | Kaggle