pass the dataset features to the IterableDataset.from_generator funct…

…ion (huggingface#5569) * pass the dataset features to the IterableDataset.from_generator function * reformat file --------- Co-authored-by: Bruno Hays <[email protected]>
gary149 · Feb 23, 2023 · b91070b · b91070b
1 parent 2070345
commit b91070b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
@@ -4983,7 +4983,9 @@ def to_iterable_dataset(self, num_shards: Optional[int] = 1) -> "IterableDataset
                 self.shard(num_shards=num_shards, index=shard_idx, contiguous=True) for shard_idx in range(num_shards)
             ]
         )
-        return IterableDataset.from_generator(Dataset._iter_shards, gen_kwargs={"shards": shards})
+        return IterableDataset.from_generator(
+            Dataset._iter_shards, features=self.features, gen_kwargs={"shards": shards}
+        )
 
     def _push_parquet_shards_to_hub(
         self,