diff --git a/.ci/docker/libtorch/build.sh b/.ci/docker/libtorch/build.sh index 8447eb0d833..c40896cb549 100755 --- a/.ci/docker/libtorch/build.sh +++ b/.ci/docker/libtorch/build.sh @@ -39,9 +39,13 @@ case ${DOCKER_TAG_PREFIX} in DOCKER_GPU_BUILD_ARG="" ;; rocm*) + # we want the patch version of 7.0 instead + if [[ "$GPU_ARCH_VERSION" == *"7.0"* ]]; then + GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2" + fi # we want the patch version of 6.4 instead if [[ "$GPU_ARCH_VERSION" == *"6.4"* ]]; then - GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2" + GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.4" fi BASE_TARGET=rocm GPU_IMAGE=rocm/dev-ubuntu-22.04:${GPU_ARCH_VERSION}-complete diff --git a/.ci/docker/manywheel/build.sh b/.ci/docker/manywheel/build.sh index 99f03f5c863..b4b50599730 100755 --- a/.ci/docker/manywheel/build.sh +++ b/.ci/docker/manywheel/build.sh @@ -75,9 +75,13 @@ case ${image} in DOCKERFILE_SUFFIX="_cuda_aarch64" ;; manylinux2_28-builder:rocm*) + # we want the patch version of 7.0 instead + if [[ "$GPU_ARCH_VERSION" == *"7.0"* ]]; then + GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2" + fi # we want the patch version of 6.4 instead if [[ "$GPU_ARCH_VERSION" == *"6.4"* ]]; then - GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2" + GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.4" fi TARGET=rocm_final MANY_LINUX_VERSION="2_28"