Do not consume unnecessary memory during sharding

There is no need to create a temporary list of a potentially very large
step/world size.

Author: janEbert

src/datasets/iterable_dataset.py

@@ -556,12 +556,8 @@ def _init_state_dict(self) -> dict:
 
     def __iter__(self):
         ex_iterator = iter(self.ex_iterable)
-        while True:
-            batch = list(islice(ex_iterator, self.step))
-            if len(batch) > self.offset:
-                yield batch[self.offset]
-            else:
-                break
+        for batch in islice(ex_iterator, self.offset, None, self.step):
+            yield batch
 
     def shuffle_data_sources(self, generator: np.random.Generator) -> "StepExamplesIterable":
         return StepExamplesIterable(