diff --git a/MLPerf/training/image_segmentation/pytorch/U-Net3D.def b/MLPerf/training/image_segmentation/pytorch/U-Net3D.def
new file mode 100644
index 0000000000000000000000000000000000000000..8771fe0db7b856d5a056f0409caec16acf8839a3
--- /dev/null
+++ b/MLPerf/training/image_segmentation/pytorch/U-Net3D.def
@@ -0,0 +1,25 @@
+Bootstrap: localimage
+From: /proj/nsc_testing/xuan/containers/pytorch_1.7.1-cuda11.0-cudnn8-runtime.sif
+
+
+%environment
+
+export PYTHONNOUSERSITE=True
+
+
+%post 
+
+mkdir -p /workspace/unet3d
+cd /workspace/unet3d
+
+apt-get update && \
+    apt-get upgrade -y && \
+    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
+
+pip install --upgrade pip
+pip install --disable-pip-version-check -r requirements.txt
\ No newline at end of file
diff --git a/MLPerf/training/image_segmentation/pytorch/build.txt b/MLPerf/training/image_segmentation/pytorch/build.txt
index 1a679dabd10f3fb4cb20e1b5054bd66370d207b9..fd77634d28b8eb9e2b9bede17302a09144147e58 100644
--- a/MLPerf/training/image_segmentation/pytorch/build.txt
+++ b/MLPerf/training/image_segmentation/pytorch/build.txt
@@ -1,16 +1,13 @@
-https://catalog.ngc.nvidia.com/orgs/nvidia/resources/nnunet_for_pytorch/quick-start-guide
 
-apptainer build containers/nvidia_pytorch_21.11-py3.sif  docker://nvcr.io/nvidia/pytorch:21.11-py3
+apptainer build containers/pytorch_1.7.1-cuda11.0-cudnn8-runtime.sif  docker://pytorch/pytorch:1.7.1-cuda11.0-cudnn8-runtime
 
-export APPTAINER_BINDPATH=
 
-MODULE_NAME=nnunet_for_pytorch
-MODULE_VERSION=21.11.0
-CONTAINER_DIR=/proj/nsc_testing/xuan/containers/${MODULE_NAME}_${MODULE_VERSION}.sif
-DEF_DIR=/proj/nsc_testing/xuan/berzelius-benchmarks/PyTorch/Segmentation/nnUNet/${MODULE_NAME}_${MODULE_VERSION}.def
+MODULE_NAME=U-Net3D
+# MODULE_VERSION=21.11.0
+CONTAINER_DIR=/proj/nsc_testing/xuan/containers/${MODULE_NAME}.sif
+DEF_DIR=/proj/nsc_testing/xuan/berzelius-benchmarks/MLPerf/training/image_segmentation/pytorch/$MODULE_NAME.def
 
 apptainer build $CONTAINER_DIR $DEF_DIR
 
 
-# Known issue 
-https://github.com/NVIDIA/DeepLearningExamples/issues/1113
+ 
\ No newline at end of file