From 88ad55927953527a52fdb883540c80e547035d7a Mon Sep 17 00:00:00 2001 From: Jacob Silterra Date: Tue, 15 Oct 2024 10:29:37 -0400 Subject: [PATCH] Set default number of threads to min(8, cpu count) --- sybil/model.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sybil/model.py b/sybil/model.py index 407445f..8ffc98a 100644 --- a/sybil/model.py +++ b/sybil/model.py @@ -124,7 +124,8 @@ def _torch_set_num_threads(threads) -> int: if threads < 0: return torch.get_num_threads() if threads is None or threads == 0: - threads = os.cpu_count() + # I've never seen a benefit to going higher than 8 and sometimes there is a big slowdown + threads = min(8, os.cpu_count()) torch.set_num_threads(threads) return torch.get_num_threads() @@ -315,7 +316,7 @@ def predict( return_attentions : bool If True, returns attention scores for each serie. See README for details. threads : int - Number of CPU threads to use for PyTorch inference. Default is 0 (use all available cores). + Number of CPU threads to use for PyTorch inference. Returns -------