WAR the dataloader issue with forked processes holding stale referenc…

…es (apache#19924)
Zha0q1 · Feb 24, 2021 · 0d9a7fa · 0d9a7fa
1 parent ac45c6e
commit 0d9a7fa
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
@@ -607,6 +607,10 @@ def __init__(self, dataset, batch_size=None, shuffle=False, sampler=None,
         self._num_workers = num_workers if num_workers >= 0 else 0
         self._worker_pool = None
         self._prefetch = max(0, int(prefetch) if prefetch is not None else 2 * self._num_workers)
+        nd.waitall()
+        import gc
+        gc.collect()
+        nd.waitall()
         if self._num_workers > 0:
             if self._thread_pool:
                 self._worker_pool = ThreadPool(self._num_workers,