Edit Dockerfile

Dockerfile

@@ -9,17 +9,15 @@ RUN pip3 install --upgrade pip
 RUN mamba update mamba
 
 # Python
-RUN mamba install -c conda-forge jupyter-server-proxy=4.4.0
-RUN mamba install -c conda-forge jupyterlab-git==0.50.1
-RUN mamba install -c conda-forge jupyter_contrib_nbextensions
-RUN mamba install -c conda-forge myst-nb
-RUN mamba install -c conda-forge ipympl
-RUN mamba install -c conda-forge ipywidgets
-
-# nvcc --version  # nvcc not found...
-RUN mamba install -c conda-forge cudatoolkit
-RUN CONDA_OVERRIDE_CUDA="12" mamba install -c conda-forge "jaxlib=*=*cuda*"
-RUN mamba install -c conda-forge jax
+RUN mamba install jupyter-server-proxy=4.4.0 -c conda-forge 
+RUN mamba install jupyterlab-git==0.50.1 -c conda-forge 
+RUN mamba install jupyter_contrib_nbextensions -c conda-forge 
+RUN mamba install myst-nb -c conda-forge 
+RUN mamba install ipympl -c conda-forge 
+RUN mamba install ipywidgets -c conda-forge 
+
+#RUN mamba install cudatoolkit -c conda-forge 
+RUN mamba install "jaxlib=*=*cuda*" jax cuda-nvcc -c conda-forge -c nvidia 
 
 RUN mamba install -c conda-forge keras
 ENV KERAS_BACKEND=jax