diff --git a/Dockerfile b/Dockerfile
index dca772d..9b1a903 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -60,7 +60,7 @@ RUN /bin/bash -c -o pipefail \
     "source /opt/venv/bin/activate && \
      pip install --no-cache-dir pyfftw colorcet wandb pandas plotly plumed 'numpy<2.0.0' && \
      pip install --no-cache-dir git+https://github.com/i-pi/i-pi.git@v3.0.0-beta4 && \
-     pip install --no-cache-dir torch==2.1 --index-url https://download.pytorch.org/whl/${GPU_LIBRARY} && \
+     pip install --no-cache-dir torch>=2.5 --index-url https://download.pytorch.org/whl/${GPU_LIBRARY} && \
      pip install --no-cache-dir git+https://github.com/acesuit/mace.git@v0.3.5"
 ARG DATE
 RUN /bin/bash -c -o pipefail \
diff --git a/build_containers.sh b/build_containers.sh
index 855660e..f5d370c 100644
--- a/build_containers.sh
+++ b/build_containers.sh
@@ -59,9 +59,12 @@ if [ "$psiflow" = "true" ]; then
 			-t ghcr.io/molmod/$TAG \
 			-f Dockerfile . # test
 		if [ "$build_sif" = "true" ]; then
+			export TMPDIR=$(pwd)/tmp
+			mkdir -p $TMPDIR
 			apptainer build -F $TAG.sif docker-daemon:ghcr.io/molmod/$TAG
 			apptainer push $TAG.sif oras://ghcr.io/molmod/$TAG
 			rm $TAG.sif
+			rm -rf $TMPDIR
 		fi
 	done
 fi
diff --git a/configs/hortense.yaml b/configs/hortense.yaml
index 330ea83..d6ccc68 100644
--- a/configs/hortense.yaml
+++ b/configs/hortense.yaml
@@ -1,7 +1,7 @@
 ---
 parsl_log_level: WARNING
 container_engine: 'apptainer'
-container_uri: 'oras://ghcr.io/molmod/psiflow:4.0.0-rc0_cu118'
+container_uri: 'oras://ghcr.io/molmod/psiflow:4.0.0_cu118'
 default_threads: 8
 ModelEvaluation:
   cores_per_worker: 12
diff --git a/configs/lumi.yaml b/configs/lumi.yaml
index b7102c0..b5e9a14 100644
--- a/configs/lumi.yaml
+++ b/configs/lumi.yaml
@@ -1,7 +1,7 @@
 ---
 parsl_log_level: WARNING
 container_engine: 'singularity'
-container_uri: 'oras://ghcr.io/molmod/psiflow:4.0.0-rc0_rocm5.6'
+container_uri: 'oras://ghcr.io/molmod/psiflow:4.0.0_rocm6.2'
 default_threads: 8
 CP2K:
   cores_per_worker: 32
@@ -25,7 +25,6 @@ ModelEvaluation:
     max_blocks: 10
     walltime: "01:00:00"
     scheduler_options: "#SBATCH --gres=gpu:8\n"
-    worker_init: "ml LUMI/23.09 && ml partition/G && ml rocm/5.6\n"
 ModelTraining:
   cores_per_worker: 7
   gpu: true
@@ -37,5 +36,4 @@ ModelTraining:
     cores_per_node: 56
     walltime: "01:00:00"
     scheduler_options: "#SBATCH --gres=gpu:8\n"
-    worker_init: "ml LUMI/23.09 && ml partition/G && ml rocm/5.6\n"
 ...