janky workaround to install FA2 on torch 2.5.1 base image since it ta…

…kes forever to build
axolotl-ai-cloud · Nov 7, 2024 · 4ddefad · 4ddefad
1 parent 02ce520
commit 4ddefad
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
@@ -35,3 +35,7 @@ RUN git lfs install --skip-repo && \
     pip3 install awscli && \
     # The base image ships with `pydantic==1.8.2` which is not working
     pip3 install -U --no-cache-dir pydantic==1.10.10
+
+RUN if [ "$PYTHON_VERSION" != "2.5.1" ] ; then \
+        pip3 install flash-attn==2.6.3; \
+    fi