Remove double-buffering on the MXNet side (#258)

Signed-off-by: ptredak <[email protected]>
NVIDIA · Oct 31, 2018 · 2a5e14c · 2a5e14c
1 parent db82480
commit 2a5e14c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dali/python/nvidia/dali/plugin/mxnet.py b/dali/python/nvidia/dali/plugin/mxnet.py
@@ -96,7 +96,7 @@ def __init__(self,
         for p in self._pipes:
             p.build()
         # Use double-buffering of data batches
-        self._data_batches = [[None, None] for i in range(self._num_gpus)]
+        self._data_batches = [[None] for i in range(self._num_gpus)]
         self._counter = 0
         self._current_data_batch = 0
         self.output_map = output_map
@@ -168,7 +168,7 @@ def __next__(self):
 
         copy_db_index = self._current_data_batch
         # Change index for double buffering
-        self._current_data_batch = (self._current_data_batch + 1) % 2
+        self._current_data_batch = (self._current_data_batch + 1) % 1
         self._counter += self._num_gpus * self.batch_size
 
         # padding the last batch