diff --git a/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/generate_benchmark_jobs.sh b/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/generate_benchmark_jobs.sh
index 5a04b7afa4ead6799772d514a0f4de1f81e2a7a9..fd1a27ba28dd1f5813b30e72609b63f171aad843 100644
--- a/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/generate_benchmark_jobs.sh
+++ b/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/generate_benchmark_jobs.sh
@@ -18,7 +18,7 @@ EOT
 if [ "${6}" = "thin" ]; then
     cat <<EOT >>  $SBATCH_DIR
 #SBATCH -C "thin"    
-#SBATCH --reservation=nodeimage
+#SBATCH --reservation=$GPU_RESERVATION
 EOT
 else
     cat <<EOT >>  $SBATCH_DIR
diff --git a/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/submit_benchmark_jobs.sh b/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/submit_benchmark_jobs.sh
index 729eba6c01b630f7f7b6335eec564e3763064032..0d1ef6f6519e6b7a9feb4029516ffa58f638ed00 100644
--- a/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/submit_benchmark_jobs.sh
+++ b/NVIDIA/DeepLearningExamples/PyTorch/Segmentation/nnUNet/submit_benchmark_jobs.sh
@@ -5,6 +5,7 @@ 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
+export GPU_RESERVATION=nodeimage
 mkdir -p $WORK_DIR/sbatch_out $WORK_DIR/sbatch_scripts $WORK_DIR/results
 
 benchmark_modes=("train" "predict")