From c08707a70bf3fbf5d9b590d8353fef7bbcb3a7ba Mon Sep 17 00:00:00 2001
From: Sascha Wise <me@saschawise.com>
Date: Sat, 30 Mar 2024 11:02:30 -0700
Subject: [PATCH] fix cuda builds

---
 .github/workflows/release.yml | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 7c6f70f..187d11d 100644
--- a/.github/workflows/release.yml
+++ b/.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,8 +161,13 @@ jobs:
           # Install a specific cuDNN version over the default one
           cuda_version="${{ matrix.cuda_version }}"
           cudnn_version="${{ matrix.cudnn_version }}"
-          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
+          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
+          apt-get install -y --allow-downgrades --allow-change-held-packages libcudnn8-dev libcudnn8
 
       # Prevent git from checking the repository owner and erroring with "dubious ownership"
       - run: git config --global --add safe.directory '*'