diff --git a/MLPerf/training/image_segmentation/pytorch/U-Net3D.def b/MLPerf/training/image_segmentation/pytorch/U-Net3D.def
index 59ba215c08af5b2d19dcf4a02bba30f07daa075f..3496f72acec5520201879bc2ae3a12fbc000befa 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-devel.sif
+From: /proj/nsc_testing/xuan/containers/pytorch_1.7.1-cuda11.0-cudnn8-runtime.sif
 
 
 %environment