From fa0ea57f5e083daab1eb9cda39ab53a5159b077d Mon Sep 17 00:00:00 2001
From: Jeff Daily <jeff.daily@amd.com>
Date: Tue, 24 Jun 2025 21:41:42 +0000
Subject: [PATCH] [ROCm][CD] upgrade to 6.4.1 patch release (#156636)

During https://github.com/pytorch/pytorch/pull/156112, we missed upgrading the manylinux and libtorch docker images.

Fixes #155292

Pull Request resolved: https://github.com/pytorch/pytorch/pull/156636
Approved by: https://github.com/jeffdaily

Co-authored-by: Jeff Daily <jeff.daily@amd.com>
---
 .ci/docker/common/install_rocm_magma.sh | 7 ++++++-
 .ci/docker/libtorch/build.sh            | 4 ++++
 .ci/docker/manywheel/build.sh           | 4 ++++
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/.ci/docker/common/install_rocm_magma.sh b/.ci/docker/common/install_rocm_magma.sh
index 364ee23b97e..a8d8ba00b35 100644
--- a/.ci/docker/common/install_rocm_magma.sh
+++ b/.ci/docker/common/install_rocm_magma.sh
@@ -5,7 +5,12 @@ set -eou pipefail
 
 function do_install() {
     rocm_version=$1
-    rocm_version_nodot=${1//./}
+    if [[ ${rocm_version} =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
+        # chop off any patch version
+        rocm_version="${rocm_version%.*}"
+    fi
+
+    rocm_version_nodot=${rocm_version//./}
 
     # Version 2.7.2 + ROCm related updates
     MAGMA_VERSION=a1625ff4d9bc362906bd01f805dbbe12612953f6
diff --git a/.ci/docker/libtorch/build.sh b/.ci/docker/libtorch/build.sh
index b8e5af5ce33..a2c67f0aa64 100755
--- a/.ci/docker/libtorch/build.sh
+++ b/.ci/docker/libtorch/build.sh
@@ -39,6 +39,10 @@ case ${DOCKER_TAG_PREFIX} in
         DOCKER_GPU_BUILD_ARG=""
         ;;
     rocm*)
+        # we want the patch version of 6.4 instead
+        if [[ $(ver $GPU_ARCH_VERSION) -eq $(ver 6.4) ]]; then
+            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.1"
+        fi
         BASE_TARGET=rocm
         GPU_IMAGE=rocm/dev-ubuntu-22.04:${GPU_ARCH_VERSION}-complete
         PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201"
diff --git a/.ci/docker/manywheel/build.sh b/.ci/docker/manywheel/build.sh
index b364d47c6c7..a4942a65bf5 100755
--- a/.ci/docker/manywheel/build.sh
+++ b/.ci/docker/manywheel/build.sh
@@ -75,6 +75,10 @@ case ${image} in
         DOCKERFILE_SUFFIX="_cuda_aarch64"
         ;;
     manylinux2_28-builder:rocm*)
+        # we want the patch version of 6.4 instead
+        if [[ $(ver $GPU_ARCH_VERSION) -eq $(ver 6.4) ]]; then
+            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.1"
+        fi
         TARGET=rocm_final
         MANY_LINUX_VERSION="2_28"
         DEVTOOLSET_VERSION="11"