fix: use packaging.version.Version for semver comparison in setup.py

Author: fparisio

extension_cpp/requirements.txt

@@ -1,2 +1,3 @@
 torch
 numpy
+packaging

extension_cpp/setup.py

@@ -8,6 +8,7 @@
 import glob
 
 from setuptools import find_packages, setup
+from packaging.version import Version
 
 from torch.utils.cpp_extension import (
     CppExtension,
@@ -18,10 +19,7 @@
 
 library_name = "extension_cpp"
 
-if torch.__version__ >= "2.6.0":
-    py_limited_api = True
-else:
-    py_limited_api = False
+py_limited_api = Version(torch.__version__) >= Version("2.6.0")
 
 
 def get_extensions():

extension_cpp_stable/requirements.txt

@@ -1,2 +1,3 @@
 torch
 numpy
+packaging

extension_cpp_stable/setup.py

@@ -8,6 +8,7 @@
 import glob
 
 from setuptools import find_packages, setup
+from packaging.version import Version
 
 from torch.utils.cpp_extension import (
     CppExtension,
@@ -19,10 +20,7 @@
 library_name = "extension_cpp_stable"
 
 
-if torch.__version__ >= "2.6.0":
-    py_limited_api = True
-else:
-    py_limited_api = False
+py_limited_api = Version(torch.__version__) >= Version("2.6.0")
 
 
 def get_extensions():