Added ROCM_PATH in case ROCM_ROOT is not specified.

Signed-off-by: Sigbjørn Løland Bore <[email protected]>
deepmodeling · Sep 20, 2024 · 3f98f6f · 3f98f6f
1 parent c084b20
commit 3f98f6f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/backend/read_env.py b/backend/read_env.py
@@ -60,6 +60,8 @@ def get_argument_from_env() -> Tuple[str, list, list, dict, str, str]:
         cmake_minimum_required_version = "3.21"
         cmake_args.append("-DUSE_ROCM_TOOLKIT:BOOL=TRUE")
         rocm_root = os.environ.get("ROCM_ROOT")
+        if not rocm_root:
+            rocm_root = os.environ.get("ROCM_PATH")
         if rocm_root:
             cmake_args.append(f"-DCMAKE_HIP_COMPILER_ROCM_ROOT:STRING={rocm_root}")
         hipcc_flags = os.environ.get("HIP_HIPCC_FLAGS")