diff --git a/container.def b/container.def
index 5b139690a4c9e8e15873da55cd848fac4b9c613b..a8c2a8b30a87146166f1a69569019bedbb9c9780 100644
--- a/container.def
+++ b/container.def
@@ -4,7 +4,7 @@ From: nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
 
 %files
     # Copy the env.yaml file into the container
-    requirements.txt /opt/env/requirements.txt
+    requirements.yml /opt/env/requirements.yml
 
 %post
     # Install special software or additions
@@ -27,8 +27,8 @@ From: nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
     /opt/conda/bin/conda init zsh
     /opt/conda/bin/conda init tcsh
 
-    # Install Python packages from MSOT_diffusion_env.yaml into the base environment
-    /opt/conda/bin/conda env update -n base python=3.8 --prune
+    # Install Python packages from requirements.yaml into the base environment
+    /opt/conda/bin/conda env update -n base --file /opt/env/requirements.yaml --prune
 
 %environment
     # Set some environment vars
@@ -43,8 +43,6 @@ From: nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
     source /etc/profile.d/conda.sh
     source ~/.bashrc
     exec "$@"
-    pip install "jax[cpu]===0.3.14" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --use-deprecated legacy-resolver
-    pip install -r /opt/env/requirements.txt
 
 %labels
     # Add custom metadata to the container.