FA2

app.py

@@ -2,12 +2,17 @@
 Holds the gradio app itself
 """
 
+import os
 import gradio as gr
 
 from src.train_workflow import run, DEFAULT_TRAINING_ARGS
 from src.calibration_datasets import CalibrationDataset
 
 
+# TODO: install FA2 in a better way, e.g docker img
+os.system("pip install flash-attn --no-build-isolation")
+
+
 DESCRIPTION = """
 The steps to create [medusa](https://sites.google.com/view/medusa-llm) heads are the following:
 

requirements.txt

@@ -1,2 +1 @@
 medusa-llm[train]
-flash-attn --no-build-isolation