add cu126 torchvision and torchaudio aarch64 build (pytorch#6062)

related to pytorch/pytorch#138440 cc @atalman @ptrblck @nWEIdia
atalman · Dec 17, 2024 · 6f6764c · 6f6764c
1 parent d8aa5aa
commit 6f6764c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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",