set torch version to what is installed during axolotl install (#1234)

axolotl-ai-cloud · Jan 31, 2024 · 8f2b591 · 8f2b591
1 parent 5787e1a
commit 8f2b591
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -27,6 +27,7 @@ def parse_requirements():
 
     try:
         torch_version = version("torch")
+        _install_requires.append(f"torch=={torch_version}")
         if torch_version.startswith("2.1."):
             _install_requires.pop(_install_requires.index("xformers==0.0.22"))
             _install_requires.append("xformers>=0.0.23")
@@ -50,7 +51,7 @@ def parse_requirements():
     dependency_links=dependency_links,
     extras_require={
         "flash-attn": [
-            "flash-attn==2.3.3",
+            "flash-attn==2.5.0",
         ],
         "fused-dense-lib": [
             "fused-dense-lib  @ git+https://github.com/Dao-AILab/[email protected]#subdirectory=csrc/fused_dense_lib",