diff --git a/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/generate_benchmark_jobs.sh b/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/generate_benchmark_jobs.sh
index 59afff9db6ebb31a947cc972636d70e86ee3deb9..5a04b7afa4ead6799772d514a0f4de1f81e2a7a9 100644
--- a/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/generate_benchmark_jobs.sh
+++ b/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/generate_benchmark_jobs.sh
@@ -1,13 +1,8 @@
 #!/bin/bash
 
-MODULE_NAME=nnunet_for_pytorch
-MODULE_VERSION=21.11.0
-WORK_DIR=/proj/nsc_testing/xuan/berzelius-benchmarks/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet
-CONTAINER_DIR=/proj/nsc_testing/xuan/containers/${MODULE_NAME}_${MODULE_VERSION}.sif
 SBATCH_DIR=$WORK_DIR/sbatch_scripts/benchmark_${6}_${5}_dim${1}_nodes${2}_gpus${3}_batchsize_${4}.sbatch
 SBATCH_OUT_DIR=$WORK_DIR/sbatch_out/benchmark_${6}_${5}_dim${1}_nodes${2}_gpus${3}_batchsize_${4}.out
 LOG_DIR=benchmark_${6}_${5}_dim${1}_nodes${2}_gpus${3}_batchsize_${4}_amp.json
-mkdir -p $WORK_DIR/sbatch_out $WORK_DIR/sbatch_scripts
 
 cat <<EOT >  $SBATCH_DIR
 #!/bin/bash
@@ -15,7 +10,7 @@ cat <<EOT >  $SBATCH_DIR
 #SBATCH -A nsc
 #SBATCH --nodes=${2}
 #SBATCH --gpus=${3}
-#SBATCH --time=0-0:10:00
+#SBATCH --time=0-0:20:00
 #SBATCH --output=$SBATCH_OUT_DIR
 
 EOT
diff --git a/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/submit_benchmark_jobs.sh b/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/submit_benchmark_jobs.sh
index 09198942b5ad233ac61b576ecbaae0352f4213b6..729eba6c01b630f7f7b6335eec564e3763064032 100644
--- a/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/submit_benchmark_jobs.sh
+++ b/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/submit_benchmark_jobs.sh
@@ -1,7 +1,12 @@
 #!/bin/bash
 set -e
 
-WORK_DIR=/proj/nsc_testing/xuan/berzelius-benchmarks/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet
+export MODULE_NAME=nnunet_for_pytorch
+export MODULE_VERSION=21.11.0
+export WORK_DIR=/proj/nsc_testing/xuan/berzelius-benchmarks/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet
+export CONTAINER_DIR=/proj/nsc_testing/xuan/containers/${MODULE_NAME}_${MODULE_VERSION}.sif
+mkdir -p $WORK_DIR/sbatch_out $WORK_DIR/sbatch_scripts $WORK_DIR/results
+
 benchmark_modes=("train" "predict")
 # node_types=("thin" "fat")
 node_types=("thin")