fix import multiprocessing

apache · Jul 28, 2021 · 4438dcb · 4438dcb
1 parent b60a2d2
commit 4438dcb
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py
@@ -36,6 +36,7 @@
 import shutil
 import tempfile
 import multiprocessing
+from multiprocessing.pool import ThreadPool
 import logging
 
 import tvm._ffi
@@ -702,7 +703,7 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
         The build results of these MeasureInputs.
     """
     # This pool is not doing computationally intensive work, so we can use threads
-    pool = multiprocessing.pool.ThreadPool(n_parallel)
+    pool = ThreadPool(n_parallel)
     tuple_res = pool.map(
         local_build_worker,
         [
@@ -1259,7 +1260,7 @@ def rpc_runner_run(
     """
     assert len(inputs) == len(build_results), "Measure input size should be equal to build results"
     # This pool is not doing computationally intensive work, so we can use threads
-    pool = multiprocessing.pool.ThreadPool(n_parallel)
+    pool = ThreadPool(n_parallel)
     tuple_res = pool.map(
         _rpc_run_worker,
         [