From 4438dcbd5f0197bbea444d0676636210b39c17bf Mon Sep 17 00:00:00 2001 From: shingjan Date: Wed, 28 Jul 2021 15:41:10 -0700 Subject: [PATCH] fix import multiprocessing --- python/tvm/auto_scheduler/measure.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index 8d762602bfd17..10166d6658de3 100644 --- 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, [