fix NVTE_UB_WITH_MPI read (#1194)

* fix NVTE_UB_WITH_MPI read Signed-off-by: Sangkug Lym <[email protected]> * Add default value Signed-off-by: Sangkug Lym <[email protected]> --------- Signed-off-by: Sangkug Lym <[email protected]> Co-authored-by: Kirthi Shankar Sivamani <[email protected]>
NVIDIA · Sep 25, 2024 · 209b8e5 · 209b8e5
1 parent a44cb72
commit 209b8e5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py
@@ -88,7 +88,7 @@ def setup_pytorch_extension(
     # Libraries
     library_dirs = []
     libraries = []
-    if os.getenv("NVTE_UB_WITH_MPI"):
+    if bool(int(os.getenv("NVTE_UB_WITH_MPI", 0))):
         assert (
             os.getenv("MPI_HOME") is not None
         ), "MPI_HOME must be set when compiling with NVTE_UB_WITH_MPI=1"