Update setup.py

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: