fix multi worker dataloader deadlock (#9126)

starimpact · Dec 19, 2017 · 4ef905d · 4ef905d
1 parent d7e0a4f
commit 4ef905d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/mxnet/gluon/data/dataloader.py b/python/mxnet/gluon/data/dataloader.py
@@ -220,10 +220,11 @@ def __iter__(self):
 
         for idx, batch in enumerate(self._batch_sampler):
             key_queue.put((idx, batch))
+        num_batches = idx + 1
 
         data_buffer = {}
         curr_idx = 0
-        for _ in range(len(self._batch_sampler)):
+        for _ in range(num_batches):
             idx, batch = data_queue.get()
             data_buffer[idx] = batch
             while curr_idx in data_buffer: