dockerfile updates

dfm · Nov 3, 2023 · a474653 · a474653
1 parent d99789e
commit a474653
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/.github/action/Dockerfile b/.github/action/Dockerfile
@@ -1,10 +1,10 @@
 FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
 
 RUN apt-get update && \
-    DEBIAN_FRONTEND=noninteractive apt-get install -y git python3-pip cmake
+    DEBIAN_FRONTEND=noninteractive apt-get install -y git python3-pip
 
 RUN pip install --upgrade pip && \
-    pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+    pip install --upgrade "jax[cuda11_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 COPY entrypoint.sh /entrypoint.sh