fix a wrong check on torch version

THU-MIG · Dec 28, 2021 · 4a15692 · 4a15692
1 parent 4aacd11
commit 4a15692
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchpruner/model_tools.py b/torchpruner/model_tools.py
@@ -180,6 +180,6 @@ def replace_object_by_class(
 
 def normalize_onnx_parameters(**kwargs):
     torch_version = torch.__version__.split(".")
-    if torch_version[0] > "2" or len(torch_version) > 1 and torch_version[1] >= "10":
+    if int(torch_version[0]) >= 2 or len(torch_version) > 1 and int(torch_version[1]) >= 10:
         kwargs.pop("_retain_param_name", None)
     return kwargs