From d8b7216ee4063559b377ccd3e4578a5c7da5b5dd Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Mon, 9 Feb 2026 11:51:13 -0800 Subject: [PATCH] fix: handle nvidia namespace packages where __file__ is None Signed-off-by: Santosh Bhavani --- transformer_engine/common/__init__.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 02388d2e70..40933f17a9 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -245,11 +245,13 @@ def _nvidia_cudart_include_dir() -> str: return "" # Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia" - # above doesn't through. However, they don't set "__file__" attribute. - if nvidia.__file__ is None: - return "" + # above doesn't throw. However, they don't set "__file__" attribute. + if nvidia.__file__ is not None: + nvidia_root = Path(nvidia.__file__).parent + else: + nvidia_root = Path(nvidia.__path__[0]) # namespace package - include_dir = Path(nvidia.__file__).parent / "cuda_runtime" + include_dir = nvidia_root / "cuda_runtime" return str(include_dir) if include_dir.exists() else ""