WAR the dataloader issue with forked processes holding stale referenc…

…es (#19925)
apache · Feb 21, 2021 · fee8bf4 · fee8bf4
1 parent 2c4f5aa
commit fee8bf4
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/mxnet/gluon/data/dataloader.py b/python/mxnet/gluon/data/dataloader.py
@@ -650,6 +650,10 @@ def __init__(self, dataset, batch_size=None, shuffle=False, sampler=None,
                 pin_device_id=self._pin_device_id,
                 prefetch=self._prefetch, **mx_iter_args)
         else:
+            nd.waitall()
+            import gc
+            gc.collect()
+            nd.waitall()
             if self._num_workers > 0:
                 if self._thread_pool:
                     self._worker_pool = ThreadPool(self._num_workers,