From 00db8eecadab7e4c624e8e88f8182b2bc9754924 Mon Sep 17 00:00:00 2001 From: Andrey Talman Date: Tue, 17 Dec 2024 13:52:40 +0000 Subject: [PATCH] [cherry-pick] add cu126 torchvision and torchaudio aarch64 build (#6075) Cherry-Pick https://github.com/pytorch/test-infra/pull/6062 to release Co-authored-by: Ting Lu <92425201+tinglvv@users.noreply.github.com> --- tools/scripts/generate_binary_build_matrix.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/scripts/generate_binary_build_matrix.py b/tools/scripts/generate_binary_build_matrix.py index d666100532..19bbbeda57 100755 --- a/tools/scripts/generate_binary_build_matrix.py +++ b/tools/scripts/generate_binary_build_matrix.py @@ -163,7 +163,7 @@ def initialize_globals(channel: str, build_python_only: bool) -> None: XPU: "pytorch/manylinux2_28-builder:xpu-2.6", # TODO: Migrate CUDA_AARCH64 image to manylinux2_28_aarch64-builder:cuda12.4 CPU_AARCH64: "pytorch/manylinux2_28_aarch64-builder:cpu-aarch64-2.6", - CUDA_AARCH64: "pytorch/manylinuxaarch64-builder:cuda12.4-2.6", + CUDA_AARCH64: "pytorch/manylinuxaarch64-builder:cuda12.6-2.6", } LIBTORCH_CONTAINER_IMAGES = { **{ @@ -191,7 +191,7 @@ def translate_desired_cuda(gpu_arch_type: str, gpu_arch_version: str) -> str: return { CPU: "cpu", CPU_AARCH64: CPU, - CUDA_AARCH64: "cu124", + CUDA_AARCH64: "cu126", CUDA: f"cu{gpu_arch_version.replace('.', '')}", ROCM: f"rocm{gpu_arch_version}", XPU: "xpu",