diff --git a/modules/rocm.py b/modules/rocm.py
index 1557a408a..4b9cef6c7 100644
--- a/modules/rocm.py
+++ b/modules/rocm.py
@@ -202,10 +202,8 @@ else:
         return version == version_torch and bool(int(os.environ.get("TORCH_BLAS_PREFER_HIPBLASLT", "1")))
 
     def get_flash_attention_command(agent: Agent):
-        if os.environ.get("FLASH_ATTENTION_USE_TRITON_ROCM", "false").lower() == "true":
-            return "pytest git+https://github.com/ROCm/flash-attention@micmelesse/upstream_pr"
         default = "git+https://github.com/ROCm/flash-attention"
-        if agent.gfx_version >= 0x1100:
+        if agent.gfx_version >= 0x1100 and os.environ.get("FLASH_ATTENTION_USE_TRITON_ROCM", "false").lower() != "true":
             default = "git+https://github.com/ROCm/flash-attention@howiejay/navi_support"
         return os.environ.get("FLASH_ATTENTION_PACKAGE", default)