From 4ddefad3821cb338e15dd6b3589251667209ef5d Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Thu, 7 Nov 2024 14:58:56 -0500
Subject: [PATCH] janky workaround to install FA2 on torch 2.5.1 base image
 since it takes forever to build

---
 docker/Dockerfile-base | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
index 3f13bba30a..405883eda6 100644
--- 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