Dockerfile

- changed version from 3.1.1 to 1.5.0
- removed pip install jax
- added mamba install -c conda-forge jax

Dockerfile

-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