Adjust Version check

Signed-off-by: Andrea Fasoli <[email protected]>

Author: andrea-fasoli

fms_mo/aiu_addons/gptq/gptq_aiu_op.py

@@ -42,7 +42,7 @@ def register_op_decorator(pt_ver, op_namespace_id):
     """Version-dependent decorator for custom op registration."""
 
     def decorator(func):
-        if pt_ver <= Version("2.4"):
+        if pt_ver < Version("2.4"):
             return torch.library.impl_abstract(op_namespace_id)(func)
         return torch.library.register_fake(op_namespace_id)(func)
 
@@ -61,7 +61,7 @@ def register_aiu_gptq_op():
         logger.warning("AIU op has already been registered")
         return
     op_namespace_id = "gptq_gemm::i4f16_fxinputs_aiu"
-    if torch_version <= Version("2.4"):
+    if torch_version < Version("2.4"):
         torch.library.define(
             op_namespace_id,
             "(Tensor x, Tensor qw, Tensor qzeros, "

fms_mo/aiu_addons/i8i8/i8i8_aiu_op.py

@@ -36,7 +36,7 @@ def implement_op_decorator(pt_ver, op_namespace_id):
     """Version-dependent decorator for custom op implementation."""
 
     def decorator(func):
-        if pt_ver <= Version("2.4"):
+        if pt_ver < Version("2.4"):
             return torch.library.impl(op_namespace_id, "default")(func)
         return torch.library.custom_op(op_namespace_id, mutates_args=())(func)
 
@@ -47,7 +47,7 @@ def register_op_decorator(pt_ver, op_namespace_id):
     """Version-dependent decorator for custom op registration."""
 
     def decorator(func):
-        if pt_ver <= Version("2.4"):
+        if pt_ver < Version("2.4"):
             return torch.library.impl_abstract(op_namespace_id)(func)
         return torch.library.register_fake(op_namespace_id)(func)