fix the total number of returned lines in setup_iter func in RawTextIterableDataset

test/data/test_builtin_datasets.py

@@ -138,6 +138,12 @@ def test_text_classification(self):
         self._helper_test_func(len(test_iter), 7600, next(iter(test_iter))[1][:25], 'Fears for T N pension aft')
         del train_iter, test_iter
 
+    def test_num_lines_of_setup_iter_dataset(self):
+        train_iter, test_iter = torchtext.experimental.datasets.raw.AG_NEWS()
+        train_iter.setup_iter(start=10, num_lines=100)
+        _data = [item for item in train_iter]
+        self.assertEqual(len(_data), 100)
+
     def test_imdb(self):
         from torchtext.experimental.datasets import IMDB
         from torchtext.vocab import Vocab

torchtext/experimental/datasets/raw/common.py

@@ -40,7 +40,7 @@ def __iter__(self):
         for i, item in enumerate(self._iterator):
             if i < self.start:
                 continue
-            if self.num_lines and i > (self.start + self.num_lines):
+            if self.num_lines and i >= (self.start + self.num_lines):
                 break
             yield item