fix cuda builds

.github/workflows/release.yml

@@ -127,13 +127,17 @@ jobs:
           - container: nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
             xla_target: cuda118
             cuda_version: "11.8.0"
+            cuda_version_short: "11.8"
             cudnn_version: "8.6.0"
             python: "3.11"
+            cudnn_long_version: "8.6.0.163"
           - container: nvidia/cuda:12.0.0-cudnn8-devel-ubuntu20.04
             xla_target: cuda120
             cuda_version: "12.0.0"
+            cuda_version_short: "12.0"
             cudnn_version: "8.8.0"
             python: "3.11"
+            cudnn_long_version: "8.8.0.121"
     container: ${{ matrix.container }}
     env:
       # Set the missing UTF-8 locale, otherwise Elixir warns
@@ -157,9 +161,16 @@ jobs:
           # Install a specific cuDNN version over the default one
           cuda_version="${{ matrix.cuda_version }}"
           cudnn_version="${{ matrix.cudnn_version }}"
+          cuda_version_short="${{ matrix.cuda_version_short }}"
+          cudnn_long_version="${{matrix.cudnn_long_version}}"
+          wget "https://developer.download.nvidia.com/compute/redist/cudnn/v${cudnn_version}/local_installers/${cuda_version_short}/cudnn-local-repo-ubuntu2204-${cudnn_long_version}_1.0-1_amd64.deb"
+          dpkg -i "cudnn-local-repo-ubuntu2204-${cudnn_long_version}_1.0-1_amd64.deb"
+          cp /var/cudnn-local-repo-ubuntu2204-${cudnn_long_version}/*.gpg /usr/share/keyrings/
+          apt-get update
           cudnn_package_version="$(apt-cache madison libcudnn8 | grep "${cudnn_version}.*-1+cuda${cuda_version%.*}")"
           apt-get install -y --allow-downgrades --allow-change-held-packages libcudnn8=$cudnn_package_version libcudnn8-dev=$cudnn_package_version
 
+
       # Prevent git from checking the repository owner and erroring with "dubious ownership"
       - run: git config --global --add safe.directory '*'
       # Proceed with the regular steps