add more cuda version for aarch64

Author: lanluo-nvidia

.github/scripts/install-torch-tensorrt.sh

@@ -5,17 +5,18 @@ INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION}
 PLATFORM=$(python -c "import sys; print(sys.platform)")
 
 if [[ $(uname -m) == "aarch64" ]]; then
-    # install cuda 12.8 only for aarch64 for now
-    VER="12-8"
+    # install cuda for aarch64
+    # CU_VERSION: cu128 --> CU_VER: 12-8
+    CU_VER=${CU_VERSION:2:2}-${CU_VERSION:4:1}
     dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo
-    dnf -y install cuda-compiler-${VER}.aarch64 \
-                cuda-libraries-${VER}.aarch64 \
-                cuda-libraries-devel-${VER}.aarch64
+    dnf -y install cuda-compiler-${CU_VER}.aarch64 \
+                   cuda-libraries-${CU_VER}.aarch64 \
+                   cuda-libraries-devel-${CU_VERER}.aarch64
     dnf clean all
     export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
     ls -lart /usr/local/
     nvcc --version
-    echo "cuda installed successfully"
+    echo "cuda ${CU_VER} installed successfully"
 fi
 
 # Install all the dependencies required for Torch-TensorRT

packaging/pre_build_script.sh

@@ -15,15 +15,16 @@ if [[ $(uname -m) == "aarch64" ]]; then
     # aarch64 does not have envsubst pre-installed in the image, install it here
     curl -L  https://github.com/a8m/envsubst/releases/download/v1.4.2/envsubst-Linux-arm64 -o envsubst \
     && mv envsubst /usr/bin/envsubst && chmod +x /usr/bin/envsubst
-    # install cuda 12.8 only for aarch64 for now
-    VER="12-8"
+    # install cuda for aarch64
+    # CU_VERSION: cu128 --> CU_VER: 12-8
+    CU_VER=${CU_VERSION:2:2}-${CU_VERSION:4:1}
     dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo
-    dnf -y install cuda-compiler-${VER}.aarch64 \
-                   cuda-libraries-${VER}.aarch64 \
-                   cuda-libraries-devel-${VER}.aarch64
+    dnf -y install cuda-compiler-${CU_VER}.aarch64 \
+                   cuda-libraries-${CU_VER}.aarch64 \
+                   cuda-libraries-devel-${CU_VERER}.aarch64
     dnf clean all
     export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
-    echo "cuda installed successfully"
+    echo "cuda ${CU_VER} installed successfully"
 fi
 
 curl -L https://github.com/bazelbuild/bazelisk/releases/download/v1.26.0/bazelisk-linux-${BAZEL_PLATFORM} \