fix issue 3259 (#3260)

Author: lanluo-nvidia

py/torch_tensorrt/_compile.py

@@ -343,6 +343,15 @@ def convert_method_to_trt_engine(
         enabled_precisions if enabled_precisions is not None else {torch.float}
     )
 
+    if not arg_inputs and not inputs:
+        raise AssertionError("'arg_inputs' and 'inputs' should not both be None.")
+
+    elif arg_inputs and inputs:
+        raise AssertionError(
+            "'arg_inputs' and 'inputs' should not be used at the same time."
+        )
+    arg_inputs = arg_inputs or inputs
+
     module_type = _parse_module_type(module)
     target_ir = _get_target_fe(module_type, ir)
     if target_ir == _IRType.ts:
@@ -366,15 +375,6 @@ def convert_method_to_trt_engine(
         )
     elif target_ir == _IRType.dynamo:
         # Prepare torch and torchtrt inputs
-        if not arg_inputs and not inputs:
-            raise AssertionError("'arg_inputs' and 'inputs' should not both be None.")
-
-        elif arg_inputs and inputs:
-            raise AssertionError(
-                "'arg_inputs' and 'inputs' should not be used at the same time."
-            )
-        arg_inputs = arg_inputs or inputs
-
         if kwarg_inputs is None:
             kwarg_inputs = {}
 

tests/py/ts/api/test_classes.py

@@ -232,7 +232,7 @@ def test_dynamic_shape(self):
 )
 class TestTorchTensorRTModule(unittest.TestCase):
     @staticmethod
-    def _get_trt_mod():
+    def _get_trt_mod(via_ts: bool = False):
         class Test(torch.nn.Module):
             def __init__(self):
                 super(Test, self).__init__()
@@ -244,9 +244,14 @@ def forward(self, x):
                 return out
 
         mod = torch.jit.script(Test())
-        test_mod_engine_str = torchtrt.ts.convert_method_to_trt_engine(
-            mod, "forward", inputs=[torchtrt.Input((2, 10))]
-        )
+        if via_ts:
+            test_mod_engine_str = torchtrt.ts.convert_method_to_trt_engine(
+                mod, "forward", inputs=[torchtrt.Input((2, 10))]
+            )
+        else:
+            test_mod_engine_str = torchtrt.convert_method_to_trt_engine(
+                mod, "forward", inputs=[torchtrt.Input((2, 10))]
+            )
         return TorchTensorRTModule(
             name="test",
             serialized_engine=test_mod_engine_str,
@@ -301,9 +306,12 @@ def forward(self, x):
             )
 
     def test_set_get_profile_path_prefix(self):
-        trt_mod = TestTorchTensorRTModule._get_trt_mod()
-        trt_mod.engine.profile_path_prefix = "/tmp/"
-        self.assertTrue(trt_mod.engine.profile_path_prefix == "/tmp/")
+        for trt_mod in (
+            TestTorchTensorRTModule._get_trt_mod(),
+            TestTorchTensorRTModule._get_trt_mod(via_ts=True),
+        ):
+            trt_mod.engine.profile_path_prefix = "/tmp/"
+            self.assertTrue(trt_mod.engine.profile_path_prefix == "/tmp/")
 
     def test_get_layer_info(self):
         """
@@ -321,11 +329,14 @@ def test_get_layer_info(self):
 
         import json
 
-        trt_mod = TestTorchTensorRTModule._get_trt_mod()
-        trt_json = json.loads(trt_mod.get_layer_info())
-        [self.assertTrue(k in trt_json.keys()) for k in ["Layers", "Bindings"]]
-        self.assertTrue(len(trt_json["Layers"]) == 2)
-        self.assertTrue(len(trt_json["Bindings"]) == 2)
+        for trt_mod in (
+            TestTorchTensorRTModule._get_trt_mod(),
+            TestTorchTensorRTModule._get_trt_mod(via_ts=True),
+        ):
+            trt_json = json.loads(trt_mod.get_layer_info())
+            [self.assertTrue(k in trt_json.keys()) for k in ["Layers", "Bindings"]]
+            self.assertTrue(len(trt_json["Layers"]) == 2)
+            self.assertTrue(len(trt_json["Bindings"]) == 2)
 
 
 if __name__ == "__main__":