able to turn off nim even if available

lucidrains · Sep 22, 2024 · 6931bdc · 6931bdc
1 parent d158479
commit 6931bdc
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/.env.sample b/.env.sample
@@ -1,3 +1,4 @@
 TYPECHECK=True
 DEBUG=False
 DEEPSPEED_CHECKPOINTING=False
+USE_NIM=False
diff --git a/alphafold3_pytorch/tensor_typing.py b/alphafold3_pytorch/tensor_typing.py
@@ -101,6 +101,9 @@ def package_available(package_name: str) -> bool:
     HAS_NIM = False
     NIM_VERSION = None
 
+USE_NIM = env.bool('USE_NIM', HAS_NIM)
+
+assert not (USE_NIM and not HAS_NIM), 'you cannot use Nim if it is not available'
 assert not HAS_NIM or version.parse(NIM_VERSION) >= version.parse('2.0.8'), 'nim version must be 2.0.8 or above'
 
 # check is github ci
@@ -136,5 +139,6 @@ def package_available(package_name: str) -> bool:
     beartype_isinstance,
     checkpoint,
     IS_DEBUGGING,
-    IS_GITHUB_CI
+    IS_GITHUB_CI,
+    USE_NIM
 ]