diff --git a/docker/flexflow-environment/Dockerfile b/docker/flexflow-environment/Dockerfile
index 596d099f79..d571befdda 100644
--- a/docker/flexflow-environment/Dockerfile
+++ b/docker/flexflow-environment/Dockerfile
@@ -114,12 +114,12 @@ COPY docker/flexflow-environment/install_pytorch.sh /usr/local/bin/install_pytor
 RUN if [ "$FF_GPU_BACKEND" == "cuda" ] ; then \
         echo "FF_GPU_BACKEND: ${FF_GPU_BACKEND}. Installing PyTorch with CUDA"; \
         chmod +x /usr/local/bin/install_pytorch.sh && \
-        /usr/local/bin/install_pytorch.sh ${cuda_version} && \
-        rm /usr/local/bin/install_pytorch.sh; \
+        /usr/local/bin/install_pytorch.sh ${cuda_version}; \
     else \
         echo "FF_GPU_BACKEND: ${FF_GPU_BACKEND}. Installing CPU-only PyTorch"; \
         pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu; \
     fi
+RUN rm /usr/local/bin/install_pytorch.sh
 # Various dependencies
 RUN pip3 install transformers>=4.31.0 sentencepiece einops
 RUN pip3 install tensorflow notebook