Fix reporing of the end of epoch in MXNet and pyTorch plugins (#180)

Signed-off-by: ptredak <ptredak@nvidia.com>
NVIDIA · Sep 22, 2018 · 49c1e43 · 49c1e43
1 parent e811225
commit 49c1e43
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/dali/python/nvidia/dali/plugin/mxnet.py b/dali/python/nvidia/dali/plugin/mxnet.py
@@ -116,7 +116,7 @@ def __next__(self):
             batch = self._first_batch
             self._first_batch = None
             return batch
-        if self._counter > self._size:
+        if self._counter >= self._size:
             raise StopIteration
         # Gather outputs
         outputs = []
@@ -177,7 +177,7 @@ def reset(self):
         DALI iterators do not support resetting before the end of the epoch
         and will ignore such request.
         """
-        if self._counter > self._size:
+        if self._counter >= self._size:
             self._counter = self._counter % self._size
         else:
             logging.warning("DALI iterator does not support resetting while epoch is not finished. Ignoring...")

diff --git a/dali/python/nvidia/dali/plugin/pytorch.py b/dali/python/nvidia/dali/plugin/pytorch.py
@@ -97,7 +97,7 @@ def __next__(self):
             batch = self._first_batch
             self._first_batch = None
             return batch
-        if self._counter > self._size:
+        if self._counter >= self._size:
             raise StopIteration
         # Gather outputs
         outputs = []
@@ -169,7 +169,7 @@ def reset(self):
         DALI iterators do not support resetting before the end of the epoch
         and will ignore such request.
         """
-        if self._counter > self._size:
+        if self._counter >= self._size:
             self._counter = self._counter % self._size
         else:
             logging.warning("DALI iterator does not support resetting while epoch is not finished. Ignoring...")