diff --git a/PyTorch/Segmentation/nnUNet/run_benchmark_single_node.sh b/PyTorch/Segmentation/nnUNet/run_benchmark_single_node.sh
index 47b800b334110e6f91738e7b5146e279062c006a..029185764e9fe290d4d288cf2ca8adfc108b5186 100644
--- a/PyTorch/Segmentation/nnUNet/run_benchmark_single_node.sh
+++ b/PyTorch/Segmentation/nnUNet/run_benchmark_single_node.sh
@@ -1,6 +1,11 @@
 #!/bin/bash
 
-cat <<EOT > scripts/benchmark_dim${1}_nodes${2}_gpus${3}_batchsize${4}_iteration${5}.sbatch
+MODULE_NAME=nnunet_for_pytorch
+MODULE_VERSION=21.11.0
+WORK_DIR=/proj/nsc_testing/xuan/berzelius-benchmarks/PyTorch/Segmentation/nnUNet
+
+cat <<EOT >  $WORK_DIR/benchmark_dim${1}_nodes${2}_gpus${3}_batchsize${4}_iteration${5}.sbatch
+
 #!/bin/bash
 
 #SBATCH -A nsc
@@ -15,14 +20,14 @@ MODULE_VERSION=21.11.0
 WORK_DIR=/proj/nsc_testing/xuan/berzelius-benchmarks/PyTorch/Segmentation/nnUNet
 CONTAINER_DIR=/proj/nsc_testing/xuan/containers/${MODULE_NAME}_${MODULE_VERSION}.sif
 
-mkdir -p $WORK_DIR/sbatch_out $WORK_DIR/benchmark_results
+mkdir -p $WORK_DIR/sbatch_out 
 
 cd $WORK_DIR
-rm -f benchmark_results/benchmark_dim${1}_nodes${2}_gpus${3}_batchsize${4}_tf32_iteration${5}.json
-srun apptainer exec --nv -B ${PWD}/data:/data -B ${PWD}/results:/results nvidia_nnu-net_for_pytorch.sif bash -c "cd /workspace/nnunet_pyt && python scripts/benchmark.py --mode train --gpus ${3} --dim ${1} --batch_size ${4} --nodes ${2} --logname='benchmark_dim${1}_nodes${2}_gpus${3}_batchsize${4}_tf32_iteration${5}.json'"
+rm -f results/benchmark_dim${1}_nodes${2}_gpus${3}_batchsize${4}_tf32_iteration${5}.json
+srun apptainer exec --nv -B ${PWD}/data:/data -B ${PWD}/results:/results $CONTAINER_DIR bash -c "cd /workspace/nnunet_pyt && python scripts/benchmark.py --mode train --gpus ${3} --dim ${1} --batch_size ${4} --nodes ${2} --logname='benchmark_dim${1}_nodes${2}_gpus${3}_batchsize${4}_tf32_iteration${5}.json'"
 
-rm -f benchmark_results/benchmark_dim${1}_nodes${2}_gpus${3}_batchsize${4}_amp_iteration${5}.json
-srun apptainer exec --nv -B ${PWD}/data:/data -B ${PWD}/results:/results nvidia_nnu-net_for_pytorch.sif bash -c "cd /workspace/nnunet_pyt && python scripts/benchmark.py --mode train --gpus ${3} --dim ${1} --batch_size ${4} --nodes ${2} --amp --logname='benchmark_dim${1}_nodes${2}_gpus${3}_batchsize${4}_amp_iteration${5}.json'" 
+rm -f results/benchmark_dim${1}_nodes${2}_gpus${3}_batchsize${4}_amp_iteration${5}.json
+srun apptainer exec --nv -B ${PWD}/data:/data -B ${PWD}/results:/results $CONTAINER_DIR bash -c "cd /workspace/nnunet_pyt && python scripts/benchmark.py --mode train --gpus ${3} --dim ${1} --batch_size ${4} --nodes ${2} --amp --logname='benchmark_dim${1}_nodes${2}_gpus${3}_batchsize${4}_amp_iteration${5}.json'" 
 
 
 EOT