forgot flash-attention in check

axolotl-ai-cloud · Jan 10, 2024 · f2a975d · f2a975d
1 parent a068283
commit f2a975d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -12,7 +12,10 @@ def parse_requirements():
         lines = [r.strip() for r in requirements_file.readlines()]
         for line in lines:
             is_extras = (
-                "flash-attn" in line or "deepspeed" in line or "mamba-ssm" in line
+                "flash-attn" in line
+                or "flash-attention" in line
+                or "deepspeed" in line
+                or "mamba-ssm" in line
             )
             if line.startswith("--extra-index-url"):
                 # Handle custom index URLs