Pin FA version (#254)

pin FA version Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Pin FA version (#254)
pin FA version Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
5495883c · Kirthi Shankar Sivamani · GitHub · 37bbfc76 · 5495883c
Unverified Commit 5495883c authored Jun 01, 2023 by Kirthi Shankar Sivamani Committed by GitHub Jun 01, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

setup.py setup.py +1 -1

No files found.
--- a/setup.py
+++ b/setup.py
@@ -283,7 +283,7 @@ def setup_requirements() -> Tuple[List[str], List[str], List[str]]:
    # Framework-specific requirements
    if "pytorch" in frameworks():
-        add_unique(install_reqs, ["torch", "flash-attn>=1.0.2"])
+        add_unique(install_reqs, ["torch", "flash-attn==1.0.6"])
        add_unique(test_reqs, ["numpy", "onnxruntime", "torchvision"])
    if "jax" in frameworks():
        if not found_pybind11():