From 4403d06eca2c8deced68903905708f0698644bc6 Mon Sep 17 00:00:00 2001
From: Niklas Muennighoff <n.muennighoff@gmail.com>
Date: Mon, 19 Dec 2022 01:40:21 +0400
Subject: [PATCH] Simplify skipping (#5373)

Update iterable_dataset.py
---
 src/datasets/iterable_dataset.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py
index 3234a3df744..09becb7f3f3 100644
--- a/src/datasets/iterable_dataset.py
+++ b/src/datasets/iterable_dataset.py
@@ -608,10 +608,7 @@ def __init__(self, ex_iterable: _BaseExamplesIterable, n: int):
         self.n = n
 
     def __iter__(self):
-        ex_iterator = iter(self.ex_iterable)
-        for _ in islice(ex_iterator, self.n):
-            pass
-        yield from ex_iterator
+        yield from islice(self.ex_iterable, self.n, None)
 
     def shuffle_data_sources(self, generator: np.random.Generator) -> "SkipExamplesIterable":
         """Doesn't shuffle the wrapped examples iterable since it would skip examples from other shards instead."""