From 4ef905d8367fd838696b30b677cc03ed28de8692 Mon Sep 17 00:00:00 2001 From: Eric Junyuan Xie Date: Mon, 18 Dec 2017 21:36:03 -0800 Subject: [PATCH] fix multi worker dataloader deadlock (#9126) --- python/mxnet/gluon/data/dataloader.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/mxnet/gluon/data/dataloader.py b/python/mxnet/gluon/data/dataloader.py index beb228ec..48ff4ef0 100644 --- 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: