Merge pull request #578 from RasmusOrsoe/deployer_threds

Add check for torch threads in deployer
graphnet-team · Sep 5, 2023 · 7857fb8 · 7857fb8
2 parents 27fb655 + 1e791f3
commit 7857fb8
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/graphnet/deployment/i3modules/deployer.py b/src/graphnet/deployment/i3modules/deployer.py
@@ -58,8 +58,11 @@ def __init__(
         """
         # This makes sure that one worker cannot access more
         # than 1 core's worth of compute.
-        torch.set_num_threads(1)
-        torch.set_num_interop_threads(1)
+
+        if torch.get_num_interop_threads() > 1:
+            torch.set_num_interop_threads(1)
+        if torch.get_num_threads() > 1:
+            torch.set_num_threads(1)
         # Check
         if isinstance(graphnet_modules, list):
             self._modules = graphnet_modules