update

flexflow · Oct 19, 2024 · 89f10f4 · 89f10f4
1 parent 4f6990f
commit 89f10f4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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