From b8733ec10d09490788ad824ac521ad49d4f31578 Mon Sep 17 00:00:00 2001
From: Benson Ma <bensonma415@meta.com>
Date: Tue, 18 Feb 2025 15:30:27 -0800
Subject: [PATCH] Nova job update (#3704)

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/792

- Update Nova jobs to account for CUDA 12.8 builds.  This appears to fix the existing nova builds for both aarch64 and x86

Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/3704

Reviewed By: sryap, spcyppt

Differential Revision: D69805713

Pulled By: q10

fbshipit-source-id: b413df76ea36c21f754371cc324f33e3d75e5551
---
 .github/scripts/nova_dir.bash | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/scripts/nova_dir.bash b/.github/scripts/nova_dir.bash
index deb534a9b1..c93df02437 100644
--- a/.github/scripts/nova_dir.bash
+++ b/.github/scripts/nova_dir.bash
@@ -18,7 +18,8 @@ if [[ "$CONDA_ENV" != "" ]]; then export CONDA_RUN="conda run --no-capture-outpu
 
 if  [[ "$CU_VERSION" == "cu121" ]] ||
     [[ "$CU_VERSION" == "cu124" ]] ||
-    [[ "$CU_VERSION" == "cu126" ]]; then
+    [[ "$CU_VERSION" == "cu126" ]] ||
+    [[ "$CU_VERSION" == "cu128" ]]; then
     export TORCH_CUDA_ARCH_LIST="7.0;8.0;9.0;9.0a"
     echo "Set TORCH_CUDA_ARCH_LIST to: ${TORCH_CUDA_ARCH_LIST}"