Update setup.py (#360)

Check if environment variable is set before retrieving variable at setup time.
torchmd · Feb 14, 2025 · 718e401 · 718e401
1 parent 6de4a2e
commit 718e401
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/setup.py b/setup.py
@@ -9,10 +9,13 @@
 
 
 # If WITH_CUDA is defined
-if os.environ.get("WITH_CUDA", "0") == "1":
-    use_cuda = True
-elif os.environ.get("WITH_CUDA", "0") == "0":
-    use_cuda = False
+if os.getenv("WITH_CUDA") is None:
+    if os.environ.get("WITH_CUDA", "0") == "1":
+        use_cuda = True
+    elif os.environ.get("WITH_CUDA", "0") == "0":
+        use_cuda = False
+    else:
+        raise ValueError("Invalid flag with WITH_CUDA environment variable. Expected '0' or '1'")
 else:
     use_cuda = torch.cuda._is_compiled()