Set OMP_NUM_THREADS=1

10 months ago · fd4d1ee8b7
parent 37153699fd
commit fd4d1ee8b7
1 changed files with 3 additions and 0 deletions
--- a/src/petals/init.py
+++ b/src/petals/init.py
@ -8,6 +8,9 @@ if platform.system() == "Darwin":
    os.environ.setdefault("no_proxy", "*")
    os.environ.setdefault("OBJC_DISABLE_INITIALIZE_FORK_SAFETY", "YES")

+    # Forward passes may freeze without that
+    os.environ.setdefault("OMP_NUM_THREADS", "1")
+
 import hivemind
 import transformers
 from packaging import version