pytorch 230

huggingface · Jun 13, 2024 · 90ee73d · 90ee73d
1 parent 08bff33
commit 90ee73d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -67,7 +67,7 @@ RUN conda create -p /app/env -y python=3.10
 
 SHELL ["conda", "run","--no-capture-output", "-p","/app/env", "/bin/bash", "-c"]
 
-RUN conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia && \
+RUN conda install pytorch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 pytorch-cuda=12.1 -c pytorch -c nvidia && \
     conda clean -ya && \
     conda install -c "nvidia/label/cuda-12.1.1" cuda-nvcc && conda clean -ya && \
     conda install xformers -c xformers && conda clean -ya
@@ -78,7 +78,7 @@ COPY --chown=1000:1000 . /app/
 RUN pip install -e . && \
     python -m nltk.downloader punkt && \
     pip install -U ninja && \
-    pip install -U https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.9.post1/flash_attn-2.5.9.post1+cu122torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl && \
+    pip install -U flash-attn --no-build-isolation && \
     pip install -U deepspeed && \
     pip install "unsloth[cu121-ampere-torch230] @ git+https://github.com/unslothai/unsloth.git" --no-deps && \
     pip cache purge