diff --git a/MLPerf/training/image_segmentation/pytorch/U-Net3D.def b/MLPerf/training/image_segmentation/pytorch/U-Net3D.def
index c145a55e5be29cfaa064ed8d1c99d4a71e0c8588..e60aa38155c3047b65416e40bd3a984fafe71ecc 100644
--- a/MLPerf/training/image_segmentation/pytorch/U-Net3D.def
+++ b/MLPerf/training/image_segmentation/pytorch/U-Net3D.def
@@ -1,5 +1,5 @@
 Bootstrap: localimage
-From: /proj/nsc_testing/xuan/containers/pytorch_1.7.1-cuda11.0-cudnn8-runtime.sif
+From: /proj/nsc_testing/xuan/containers/pytorch_1.7.1-cuda11.0-cudnn8-devel.sif
 
 
 %environment
diff --git a/MLPerf/training/image_segmentation/pytorch/benchmark.txt b/MLPerf/training/image_segmentation/pytorch/benchmark.txt
index 65953d5ef1e3eebb83eb392042e85d6f533612d5..6aeed748ac9cf27d5c0d98784bb961e0de5b2625 100644
--- a/MLPerf/training/image_segmentation/pytorch/benchmark.txt
+++ b/MLPerf/training/image_segmentation/pytorch/benchmark.txt
@@ -22,5 +22,5 @@ rm -rf kits19
 
 apptainer exec -B ${WORK_DIR}/raw-data:/raw-data -B ${WORK_DIR}/data:/data -B ${WORK_DIR}/results:/results $CONTAINER_DIR bash -c "cd /workspace/unet3d && python3 preprocess_dataset.py --data_dir /raw-data --results_dir /data"
 
-apptainer exec -B ${WORK_DIR}/raw-data:/raw-data -B ${WORK_DIR}/data:/data -B ${WORK_DIR}/results:/results $CONTAINER_DIR bash -c "bash /workspace/unet3d/run_and_time.sh 1"
+apptainer exec -B ${WORK_DIR}/raw-data:/raw-data -B ${WORK_DIR}/data:/data -B ${WORK_DIR}/results:/results $CONTAINER_DIR bash -c "cd /workspace/unet3d && bash run_and_time.sh 1"
 
diff --git a/MLPerf/training/image_segmentation/pytorch/build.txt b/MLPerf/training/image_segmentation/pytorch/build.txt
index fd77634d28b8eb9e2b9bede17302a09144147e58..59ad0a5c1e4cd7fa613054ba736eaa6205cf6ce4 100644
--- a/MLPerf/training/image_segmentation/pytorch/build.txt
+++ b/MLPerf/training/image_segmentation/pytorch/build.txt
@@ -1,5 +1,5 @@
 
-apptainer build containers/pytorch_1.7.1-cuda11.0-cudnn8-runtime.sif  docker://pytorch/pytorch:1.7.1-cuda11.0-cudnn8-runtime
+apptainer build containers/pytorch_1.7.1-cuda11.0-cudnn8-devel.sif  docker://pytorch/pytorch:1.7.1-cuda11.0-cudnn8-devel
 
 
 MODULE_NAME=U-Net3D