From 8f2b591baf7219ecc6cedb87b5c7b23464e00da7 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Wed, 31 Jan 2024 08:47:34 -0500
Subject: [PATCH] set torch version to what is installed during axolotl install
 (#1234)

---
 setup.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index b0c9ecbc3..6f816ce4a 100644
--- 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 @@ setup(
     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/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib",