WAR the dataloader issue with forked processes holding stale referenc…

…es (apache#19925)
chinakook · May 2, 2021 · 534a717 · 534a717
1 parent 572d132
commit 534a717
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
@@ -635,6 +635,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,