diff --git a/Dockerfile b/Dockerfile
index 0ae54c4d9ee579afdfd283174d2b1f5641b96db2..121584af353e65037571ab08cca4146570472f50 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM jupyter/tensorflow-notebook:hub-3.1.1
+FROM jupyter/tensorflow-notebook:hub-1.5.0
 
 USER root
 
@@ -9,6 +9,6 @@ RUN mamba update mamba
 # Conda packages
 RUN mamba install -c conda-forge control
 RUN mamba install -c conda-forge casadi
+RUN mamba install -c conda-forge jax
 
-# Install pip packages
-RUN pip install -c conda-forge jax
\ No newline at end of file
+# Install pip packages
\ No newline at end of file