Hopefully fix JAX not seeing GPUs

rainx0r · Aug 22, 2024 · c37f0cd · c37f0cd
1 parent ac76843
commit c37f0cd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/metaworld-jax/Dockerfile b/metaworld-jax/Dockerfile
@@ -6,7 +6,7 @@ RUN apt-get update && apt install -y --no-install-recommends git python3-pip lib
 WORKDIR /usr/src/app
 COPY requirements.txt ./
 RUN pip install --no-cache-dir --upgrade pip && \
-    pip install --no-cache-dir --upgrade -r requirements.txt
+    pip install --no-cache-dir --upgrade --pre -r requirements.txt
 
 ENTRYPOINT ["python"]
 
diff --git a/metaworld-jax/requirements.txt b/metaworld-jax/requirements.txt
@@ -1,5 +1,7 @@
 # Jax
-jax[cuda12]==0.4.31
+--find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
+jax[cuda12]==0.4.32.dev20240822
+
 flax==0.8.5
 distrax==0.1.5