diff --git a/docker/gpu.Dockerfile b/docker/gpu.Dockerfile
index d1258fc3..b0a977fe 100644
--- a/docker/gpu.Dockerfile
+++ b/docker/gpu.Dockerfile
@@ -29,6 +29,6 @@ ADD ./setup.py ./setup.py
 ADD ./README.md ./README.md
 ADD ./MANIFEST.in ./MANIFEST.in
 ADD ./kaggle_environments ./kaggle_environments
-RUN pip install Flask bitsandbytes accelerate vec-noise && pip install . && pytest
+RUN pip install Flask bitsandbytes accelerate vec-noise jax gymnax==0.0.8 && pip install . && pytest
 
 CMD kaggle-environments