[FIX] Fix using num_workers in omp (#7078)

Co-authored-by: zhangfucheng <zhangfucheng.jason@bytedance.com>
apache · Dec 26, 2020 · 6fcc64c · 6fcc64c
1 parent 959669c
commit 6fcc64c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/runtime/thread_pool.cc b/src/runtime/thread_pool.cc
@@ -369,8 +369,8 @@ int TVMBackendParallelLaunch(FTVMParallelLambda flambda, void* cdata, int num_ta
 #else
   int num_workers = tvm::runtime::threading::MaxConcurrency();
   if (num_task == 0) num_task = num_workers;
-  omp_set_num_threads(num_workers);
-#pragma omp parallel num_threads(num_workers)
+  omp_set_num_threads(num_task);
+#pragma omp parallel num_threads(num_task)
   {
     TVMParallelGroupEnv env;
     env.num_task = num_task;