From 73e72db00e6ab94b1b147281ff494c8e62d350d3 Mon Sep 17 00:00:00 2001
From: Maxime <672982+maximegmd@users.noreply.github.com>
Date: Wed, 7 Feb 2024 17:21:27 +0100
Subject: [PATCH] Update setup.py

---
 setup.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index f6a530dfae..d4a39b76ea 100644
--- a/setup.py
+++ b/setup.py
@@ -1,9 +1,9 @@
 """setup.py for axolotl"""
 
 import platform
+import re
 from importlib.metadata import PackageNotFoundError, version
 
-from packaging.version import Version, parse
 from setuptools import find_packages, setup
 
 
@@ -31,10 +31,20 @@ def parse_requirements():
         if "Darwin" in platform.system():
             _install_requires.pop(_install_requires.index("xformers==0.0.22"))
         else:
-            torch_version = parse(version("torch"))
+            torch_version = version("torch")
             _install_requires.append(f"torch=={torch_version}")
 
-            if torch_version >= Version("2.1"):
+            version_match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", torch_version)
+            if version_match:
+                major, minor, patch = version_match.groups()
+                major, minor = int(major), int(minor)
+                patch = (
+                    int(patch) if patch is not None else 0
+                )  # Default patch to 0 if not present
+            else:
+                raise ValueError("Invalid version format")
+
+            if (major, minor) >= (2, 1):
                 _install_requires.pop(_install_requires.index("xformers==0.0.22"))
                 _install_requires.append("xformers>=0.0.23")
     except PackageNotFoundError: