Skip to content

Commit 971ac9d

Browse files
authored
Fix aarch64 cpu/cuda include logic (#6584)
Make sure `with-cuda` and `with-cpu` flags are correctly generating to aarch64 builds Test: ``` python generate_binary_build_matrix.py --with-cpu disable --with-cuda enable --with-rocm disable --operating-system linux-aarch64 {"include": [{"python_version": "3.9", "gpu_arch_type": "cuda-aarch64", "gpu_arch_version": "12.8-aarch64", "desired_cuda": "cu128", "container_image": "pytorch/manylinuxaarch64-builder:cuda12.8", "package_type": "wheel", "build_name": "wheel-py3_9-cuda-aarch6412_8-aarch64", "validation_runner": "linux.arm64.m7g.4xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.7.0", "use_split_build": false}]} ```
1 parent 36b2e38 commit 971ac9d

File tree

1 file changed

+3
-6
lines changed

1 file changed

+3
-6
lines changed

tools/scripts/generate_binary_build_matrix.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -443,17 +443,14 @@ def generate_wheels_matrix(
443443
arches = []
444444

445445
if with_cpu == ENABLE:
446-
arches += [CPU]
447-
448-
if os == LINUX_AARCH64:
449-
# Only want the one arch as the CPU type is different and
450-
# uses different build/test scripts
451-
arches = [CPU_AARCH64] + CUDA_AARCH64_ARCHES
446+
arches += [CPU_AARCH64] if os == LINUX_AARCH64 else [CPU]
452447

453448
if with_cuda == ENABLE:
454449
upload_to_base_bucket = "no"
455450
if os in (LINUX, WINDOWS):
456451
arches += CUDA_ARCHES
452+
elif os == LINUX_AARCH64:
453+
arches += CUDA_AARCH64_ARCHES
457454

458455
if with_rocm == ENABLE and os == LINUX:
459456
arches += ROCM_ARCHES

0 commit comments

Comments
 (0)