From 209b8e5a801a7747adfa42235aedeb6bfd805874 Mon Sep 17 00:00:00 2001
From: Sangkug Lym <slym@nvidia.com>
Date: Tue, 24 Sep 2024 21:24:28 -0700
Subject: [PATCH] fix NVTE_UB_WITH_MPI read (#1194)

* fix NVTE_UB_WITH_MPI read

Signed-off-by: Sangkug Lym <slym@nvidia.com>

* Add default value

Signed-off-by: Sangkug Lym <slym@nvidia.com>

---------

Signed-off-by: Sangkug Lym <slym@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---
 build_tools/pytorch.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py
index 7aff2db324..9152229d2f 100644
--- 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"