Update Dockerfile

huggingface · Jul 28, 2024 · fa96d64 · fa96d64
1 parent 3ea1cb3
commit fa96d64
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -72,11 +72,10 @@ RUN conda install pytorch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 pytorch-c
     conda install -c "nvidia/label/cuda-12.1.1" cuda-nvcc && conda clean -ya && \
     conda install xformers -c xformers && conda clean -ya
 
-RUN pip install torch_xla~=2.3.0
-
 COPY --chown=1000:1000 . /app/
 
 RUN pip install -e . && \
+    pip install torch_xla~=2.3.0 && \
     python -m nltk.downloader punkt && \
     pip install -U ninja && \
     pip install -U flash-attn --no-build-isolation && \