add version utils to enable valid version str comparison

Author: jaybdub

tests/feature_tests/test_version_utils.py

@@ -0,0 +1,15 @@
+import pytest
+
+import torch2trt.version_utils
+
+def test_version_utils():
+    
+    a = torch2trt.version_utils.Version("10.1")
+
+    assert a >= "10.1"
+    assert a >= "10.0"
+    assert a > "7.0"
+    assert a < "11.0"
+    assert a == "10.1"
+    assert a <= "10.1"
+    assert a <= "10.2"

torch2trt/torch2trt.py

@@ -19,17 +19,10 @@
 
 from .flattener import Flattener
 from .flatten_module import Flatten, Unflatten
+from .version_utils import trt_version, torch_version
 # UTILITY FUNCTIONS
 
 
-def trt_version():
-    return trt.__version__
-
-
-def torch_version():
-    return torch.__version__
-
-
 def torch_dtype_to_trt(dtype):
     if trt_version() >= '7.0' and dtype == torch.bool:
         return trt.bool

torch2trt/version_utils.py

@@ -0,0 +1,40 @@
+import packaging.version
+import tensorrt as trt
+import torch
+
+
+def trt_version():
+    return trt.__version__
+
+
+def torch_version():
+    return torch.__version__
+
+
+class Version(packaging.version.Version):
+
+    def __ge__(self, other):
+        if isinstance(other, str):
+            other = Version(other)
+        return super().__ge__(other)
+
+    def __le__(self, other):
+        if isinstance(other, str):
+            other = Version(other)
+        return super().__le__(other)
+
+    def __eq__(self, other):
+        if isinstance(other, str):
+            other = Version(other)
+        return super().__eq__(other)
+
+    def __gt__(self, other):
+        if isinstance(other, str):
+            other = Version(other)
+        return super().__gt__(other)
+    
+    def __lt__(self, other):
+        if isinstance(other, str):
+            other = Version(other)
+        return super().__lt__(other)
+