Update smoke_test.py

juliagmt-google · Aug 19, 2024 · e18dd8f · e18dd8f
1 parent 5ac5836
commit e18dd8f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/smoke_test/smoke_test.py b/test/smoke_test/smoke_test.py
@@ -12,9 +12,9 @@
 from pathlib import Path
 
 if "MATRIX_GPU_ARCH_VERSION" in os.environ:
-    gpu_arch_ver = os.environ["MATRIX_GPU_ARCH_VERSION"]
+    gpu_arch_ver = os.getenv("MATRIX_GPU_ARCH_VERSION")
 else:
-    gpu_arch_ver = os.environ.get("GPU_ARCH_VERSION", None)  # Use fallback if available
+    gpu_arch_ver = os.getenv("GPU_ARCH_VERSION")  # Use fallback if available
 channel = os.getenv("MATRIX_CHANNEL")
 package_type = os.getenv("MATRIX_PACKAGE_TYPE")
 target_os = os.getenv("TARGET_OS", sys.platform)