Install CUDA 12.6 compat package

jax-ml · Sep 24, 2024 · 034a3a1 · 034a3a1
1 parent 12378f4
commit 034a3a1
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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]"