diff --git a/MLPerf/training/image_segmentation/pytorch/U-Net3D.def b/MLPerf/training/image_segmentation/pytorch/U-Net3D.def
index 2ccbbac65052924f53ee614af7a42b5ac26837e5..c145a55e5be29cfaa064ed8d1c99d4a71e0c8588 100644
--- a/MLPerf/training/image_segmentation/pytorch/U-Net3D.def
+++ b/MLPerf/training/image_segmentation/pytorch/U-Net3D.def
@@ -17,9 +17,9 @@ apt-get update && \
     apt-get install -y git
 apt-get install -y vim
 
-git clone https://github.com/mlperf/training.git
-mv training/image_segmentation/pytorch/*  /workspace/unet3d/
-rm -rf training
+git clone https://github.com/xuagu37/mlperf-training.git
+mv mlperf-training/image_segmentation/pytorch/*  /workspace/unet3d/
+rm -rf mlperf-training
 
 pip install --upgrade pip
 pip install --disable-pip-version-check -r requirements.txt
diff --git a/MLPerf/training/image_segmentation/pytorch/benchmark.txt b/MLPerf/training/image_segmentation/pytorch/benchmark.txt
index bdf3d4404651075b30eb27736525f3127d5025b3..65953d5ef1e3eebb83eb392042e85d6f533612d5 100644
--- a/MLPerf/training/image_segmentation/pytorch/benchmark.txt
+++ b/MLPerf/training/image_segmentation/pytorch/benchmark.txt
@@ -16,7 +16,11 @@ To download the data please follow the instructions:
 
 git clone https://github.com/neheller/kits19
 apptainer exec $CONTAINER_DIR bash -c "cd kits19 && python3 -m starter_code.get_imaging"
+mv kits19/data/* raw-data/
+rm -rf kits19
 
--B ${WORK_DIR}/data:/data -B ${WORK_DIR}/results:/results 
 
--B ${WORK_DIR}/raw-data:/raw-data 
\ No newline at end of file
+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"
+