Support arm64 for CUDA feature

src/nvidia-cuda/install.sh

@@ -44,9 +44,24 @@ export DEBIAN_FRONTEND=noninteractive
 
 check_packages wget ca-certificates
 
+# Determine system architecture and set NVIDIA repository URL accordingly
+ARCH=$(uname -m)
+case $ARCH in
+    x86_64)
+        NVIDIA_ARCH="x86_64"
+        ;;
+    aarch64 | arm64)
+        NVIDIA_ARCH="arm64"
+        ;;
+    *)
+        echo "Unsupported architecture: $ARCH"
+        exit 1
+        ;;
+esac
+
 # Add NVIDIA's package repository to apt so that we can download packages
 # Always use the ubuntu2004 repo because the other repos (e.g., debian11) are missing packages
-NVIDIA_REPO_URL="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64"
+NVIDIA_REPO_URL="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/$NVIDIA_ARCH"
 KEYRING_PACKAGE="cuda-keyring_1.0-1_all.deb"
 KEYRING_PACKAGE_URL="$NVIDIA_REPO_URL/$KEYRING_PACKAGE"
 KEYRING_PACKAGE_PATH="$(mktemp -d)"