From 034a3a1babc25151fd526d27ba733d73efd8f4be Mon Sep 17 00:00:00 2001
From: Michael Hudgins <michaelhudgins@google.com>
Date: Tue, 24 Sep 2024 14:46:03 +0000
Subject: [PATCH] Install CUDA 12.6 compat package

---
 .github/workflows/ci.yml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index e03a336f..8027f6b4 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -28,6 +28,11 @@ jobs:
       image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11
     steps:
     - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # ratchet:actions/checkout@v4
+    -   name: Setup Compat Driver
+        run: |
+          # This container should already have the CUDA apt repos setup 
+          apt-get update
+          apt-get install -y --no-install-recommends cuda-compat-12-6
     -   name: Setup Released JAX
         run: |
           pip install -U "jax[cuda12]"