diff --git a/orttraining/tools/amdgpu/Dockerfile.rocm4.1.pytorch b/orttraining/tools/amdgpu/Dockerfile.rocm4.1.pytorch
index 4920fb3e10..0a10760274 100644
--- a/orttraining/tools/amdgpu/Dockerfile.rocm4.1.pytorch
+++ b/orttraining/tools/amdgpu/Dockerfile.rocm4.1.pytorch
@@ -73,7 +73,9 @@ RUN git clone https://github.com/openucx/ucx.git \
   && cd build \
   && ../contrib/configure-opt --prefix=$UCX_DIR --without-rocm --without-knem --without-cuda \
   && make -j"$(nproc)" \
-  && make install
+  && make install \
+  && cd .. \
+  && rm -rf build
 
 # OpenMPI
 # note: require --enable-orterun-prefix-by-default for Azure machine learning compute
@@ -93,6 +95,8 @@ RUN git clone --recursive https://github.com/open-mpi/ompi.git \
                   --enable-mca-no-build=btl-uct --disable-mpi-fortran \
   && make -j"$(nproc)" \
   && make install \
+  && cd .. \
+  && rm -rf build \
   && ldconfig \
   && test -f ${OPENMPI_DIR}/bin/mpic++
 
@@ -128,6 +132,7 @@ RUN git clone --recursive https://github.com/microsoft/onnxruntime.git \
     --enable_training \
   && test -f $ORT_DIR/build/RelWithDebInfo/onnxruntime_training_bert \
   && pip install $ORT_DIR/build/RelWithDebInfo/dist/*.whl \
+  && rm -rf $ORT_DIR/build \
   && ldconfig
 
 # ONNX Runtime Training Examples