Fix bug in BufferShuffleSampler

asyml · Jul 22, 2019 · 859ec78 · 859ec78
1 parent d5bf9c2
commit 859ec78
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/texar/data/data/data_iterators.py b/texar/data/data/data_iterators.py
@@ -198,7 +198,8 @@ def __init__(self, data: DataBase[Any, Example], buffer_size: int):
 
     def _iterator_given_size(self, size) -> Iterator[int]:
         if self.buffer_size >= size:
-            return iter(torch.randperm(size).tolist())
+            yield from iter(torch.randperm(size).tolist())
+            return
 
         buffer = list(range(self.buffer_size))
         for x in range(self.buffer_size, size):