janky workaround to install FA2 on torch 2.5.1 base image since it takes forever to build (#2022)

2024-11-07 17:54:29 -05:00
parent 02ce520b7e
commit 035e9f9dd7
1 changed files with 4 additions and 0 deletions
--- 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