diff options
Diffstat (limited to 'src/mlia/nn/rewrite/core/utils/parallel.py')
| -rw-r--r-- | src/mlia/nn/rewrite/core/utils/parallel.py | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/src/mlia/nn/rewrite/core/utils/parallel.py b/src/mlia/nn/rewrite/core/utils/parallel.py index 5affc03..b1a2914 100644 --- a/src/mlia/nn/rewrite/core/utils/parallel.py +++ b/src/mlia/nn/rewrite/core/utils/parallel.py @@ -3,10 +3,10 @@ import math import os from collections import defaultdict +from multiprocessing import cpu_count from multiprocessing import Pool import numpy as np -from psutil import cpu_count os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" import tensorflow as tf @@ -25,7 +25,7 @@ class ParallelTFLiteModel(TFLiteModel): self.pool = None self.filename = filename if not num_procs: - self.num_procs = cpu_count(logical=False) + self.num_procs = cpu_count() else: self.num_procs = int(num_procs) |