fix: Repair version checking system for Torch

- Address version parsing issue for NV versions of Torch
- Add specialized check for NV Torch versions such as `2.0.0.nv23.05`

Author: gs-olive

py/torch_tensorrt/__init__.py

@@ -94,7 +94,11 @@ def _find_lib(name, paths):
 
 from torch_tensorrt import fx
 
-if version.parse(torch.__version__) >= version.parse("2.1.dev"):
+if version.parse(
+    torch.__version__
+    if ".nv" not in torch.__version__
+    else torch.__version__.split(".nv")[0]
+) >= version.parse("2.1.dev"):
     from torch_tensorrt import dynamo
     from torch_tensorrt.dynamo import backend
 

py/torch_tensorrt/dynamo/__init__.py

@@ -1,6 +1,10 @@
 import torch
 from packaging import version
 
-if version.parse(torch.__version__) >= version.parse("2.1.dev"):
+if version.parse(
+    torch.__version__
+    if ".nv" not in torch.__version__
+    else torch.__version__.split(".nv")[0]
+) >= version.parse("2.1.dev"):
     from torch_tensorrt.dynamo import fx_ts_compat
     from .backend import compile

py/torch_tensorrt/fx/test/passes/test_fix_reshape_batch_dim.py

@@ -43,7 +43,11 @@ def forward(self, x, y):
     %reshape : [num_users=1] = call_function[target=torch_tensorrt.fx.tracer.acc_tracer.acc_ops.reshape](args = (), kwargs = {input: %y, acc_out_ty: ((%getitem_1, -1, 3), None, None, None, None, None, None)})
     return reshape
 """
-        if version.parse(torch.__version__) < version.parse("2.1.0.dev20230620"):
+        if version.parse(
+            torch.__version__
+            if ".nv" not in torch.__version__
+            else torch.__version__.split(".nv")[0]
+        ) < version.parse("2.1.0.dev20230620"):
             expected_graph = expected_graph.replace("num_users", "#users")
 
         assert (

py/torch_tensorrt/fx/test/passes/test_remove_duplicate_output_args.py

@@ -57,7 +57,11 @@ def is_leaf_module(self, m, qn):
     return add
 """.strip()
 
-        if version.parse(torch.__version__) < version.parse("2.1.0.dev20230620"):
+        if version.parse(
+            torch.__version__
+            if ".nv" not in torch.__version__
+            else torch.__version__.split(".nv")[0]
+        ) < version.parse("2.1.0.dev20230620"):
             ttop_graph_expected = ttop_graph_expected.replace("num_users", "#users")
 
         assert (
@@ -71,7 +75,11 @@ def is_leaf_module(self, m, qn):
     return (x,)
 """.strip()
 
-        if version.parse(torch.__version__) < version.parse("2.1.0.dev20230620"):
+        if version.parse(
+            torch.__version__
+            if ".nv" not in torch.__version__
+            else torch.__version__.split(".nv")[0]
+        ) < version.parse("2.1.0.dev20230620"):
             ttop_a_graph_expected = ttop_a_graph_expected.replace("num_users", "#users")
 
         assert (

py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py

@@ -6,7 +6,11 @@
 
 import torch
 
-if version.parse(torch.__version__) >= version.parse("2.dev"):
+if version.parse(
+    torch.__version__
+    if ".nv" not in torch.__version__
+    else torch.__version__.split(".nv")[0]
+) >= version.parse("2.dev"):
     import torch._dynamo as torchdynamo
 
 from torch.fx.passes.infra.pass_base import PassResult

py/torch_tensorrt/fx/utils.py

@@ -160,7 +160,11 @@ def nested_decorator(f: Callable):
         def function_wrapper(*args, **kwargs):
             # Parse minimum and current Torch versions
             min_version = version.parse(min_torch_version)
-            current_version = version.parse(torch.__version__)
+            current_version = version.parse(
+                torch.__version__
+                if ".nv" not in torch.__version__
+                else torch.__version__.split(".nv")[0]
+            )
 
             if current_version < min_version:
                 raise AssertionError(