set_num_threads

Co-authored-by: Denis Kuznedelev <[email protected]>
OxTapinear · Jan 12, 2024 · 8e8b3ea · 8e8b3ea
1 parent b260d7a
commit 8e8b3ea
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/main.py b/main.py
@@ -690,7 +690,7 @@ def update_outs_parallel(
         help="Skip model quantization and immediately evaluate the loaded model",
     )
 
-    torch.set_num_threads(16)
+    torch.set_num_threads(min(16, torch.get_num_threads()))
     torch.backends.cudnn.allow_tf32 = False
     torch.backends.cuda.matmul.allow_tf32 = False
 

diff --git a/notebooks/aq_simple.ipynb b/notebooks/aq_simple.ipynb
@@ -28,7 +28,7 @@
     "from src.aq import QuantizedWeight\n",
     "\n",
     "\n",
-    "torch.set_num_threads(16)\n",
+    "torch.set_num_threads(min(16, saved_num_threads))\n",
     "torch.backends.cudnn.allow_tf32 = False\n",
     "torch.backends.cuda.matmul.allow_tf32 = False\n",
     "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",