resolve comments

py/torch_tensorrt/_compile.py

@@ -503,8 +503,8 @@ def save(
             )
         else:
             if arg_inputs is not None:
-                raise ValueError(
-                    "Provided model is a torch.jit.ScriptModule, do not allow user to provide inputs or arg_inputs."
+                logger.warning(
+                    "Provided model is a torch.jit.ScriptModule, inputs or arg_inputs is not necessary during save."
                 )
             torch.jit.save(module, file_path)
     elif module_type == _ModuleType.ep:
@@ -514,8 +514,8 @@ def save(
             )
         else:
             if arg_inputs is not None:
-                raise ValueError(
-                    "Provided model is a torch.export.ExportedProgram, do not allow user to provide inputs or arg_inputs during save, it should be provided during export and compile stage"
+                logger.warning(
+                    "Provided model is a torch.export.ExportedProgram, inputs or arg_inputs is not necessary during save, it uses the inputs or arg_inputs provided during export and compile"
                 )
             torch.export.save(module, file_path)
     elif module_type == _ModuleType.fx:
@@ -531,8 +531,8 @@ def save(
                 from torch_tensorrt.dynamo._exporter import export
 
                 if arg_inputs is not None:
-                    raise ValueError(
-                        "Provided model is a torch.fx.GraphModule and retrace is False, do not allow user to provide inputs or arg_inputs."
+                    logger.warning(
+                        "Provided model is a torch.fx.GraphModule and retrace is False, inputs or arg_inputs is not necessary during save."
                     )
                 exp_program = export(module)
                 torch.export.save(exp_program, file_path)

py/torch_tensorrt/dynamo/_compiler.py

@@ -700,7 +700,7 @@ def convert_exported_program_to_serialized_trt_engine(
     CONVERTERS.set_compilation_settings(settings)
 
     try:
-        _, interpreter_result = interpret_module_to_result(
+        interpreter_result, _ = interpret_module_to_result(
             gm,
             inputs=flattened_input_list,
             arg_inputs=arg_input_list,

py/torch_tensorrt/dynamo/_refit.py

@@ -14,9 +14,7 @@
 from torch_tensorrt.dynamo import partitioning
 from torch_tensorrt.dynamo._exporter import inline_torch_modules
 from torch_tensorrt.dynamo._settings import CompilationSettings
-from torch_tensorrt.dynamo.conversion._conversion import (
-    infer_module_output_shapes_dtypes,
-)
+from torch_tensorrt.dynamo.conversion._conversion import infer_module_outputs
 from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
     DYNAMO_CONVERTERS as CONVERTERS,
 )
@@ -74,7 +72,7 @@ def construct_refit_mapping(
         "CONSTANT": (trt.IConstantLayer, [("weights", "CONSTANT")]),
     }
 
-    _, output_dtypes = infer_module_output_shapes_dtypes(
+    _, output_dtypes = infer_module_outputs(
         module,
         inputs,
         settings.device,

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -26,7 +26,7 @@
 logger = logging.getLogger(__name__)
 
 
-def infer_module_output_shapes_dtypes(
+def infer_module_outputs(
     module: torch.fx.GraphModule,
     inputs: Sequence[Input],
     device: Device,
@@ -84,7 +84,7 @@ def interpret_module_to_result(
     arg_inputs: Optional[Sequence[Input]] = None,
     kwarg_inputs: Optional[dict[str, Any]] = None,
     engine_cache: Optional[BaseEngineCache] = None,
-) -> Tuple[List[Tuple[int]], TRTInterpreterResult]:
+) -> Tuple[TRTInterpreterResult, List[Tuple[int]]]:
     """Interpret an FX module to the output shapes and a TRTInterpreterResult
     Args:
         module: FX GraphModule to interpret
@@ -95,10 +95,10 @@ def interpret_module_to_result(
         settings: Compilation settings
         engine_cache: Engine cache instance
     Returns:
-        (List[Tuple[int]], TRTInterpreterResult)
+        (TRTInterpreterResult, List[Tuple[int]])
     """
     if arg_inputs is not None:
-        output_shapes, output_dtypes = infer_module_output_shapes_dtypes(
+        output_shapes, output_dtypes = infer_module_outputs(
             module,
             arg_inputs,
             settings.device,
@@ -107,7 +107,7 @@ def interpret_module_to_result(
         )
     else:
         # args and kwargs are combined and flattened to one list
-        output_shapes, output_dtypes = infer_module_output_shapes_dtypes(
+        output_shapes, output_dtypes = infer_module_outputs(
             module,
             inputs,
             settings.device,
@@ -145,7 +145,7 @@ def convert_module(
         PythonTorchTensorRTModule or TorchTensorRTModule
     """
 
-    output_shapes, interpreter_result = interpret_module_to_result(
+    interpreter_result, output_shapes = interpret_module_to_result(
         module, inputs, settings, engine_cache=engine_cache
     )
 

tests/py/dynamo/conversion/harness.py

@@ -16,9 +16,7 @@
 
 # Use interpreter, input spec, and test case from fx_ts_compat to test Dynamo Converter Registry
 from torch_tensorrt.dynamo.conversion import TRTInterpreter
-from torch_tensorrt.dynamo.conversion._conversion import (
-    infer_module_output_shapes_dtypes,
-)
+from torch_tensorrt.dynamo.conversion._conversion import infer_module_outputs
 from torch_tensorrt.dynamo.lowering import (
     get_decompositions,
     post_lowering,
@@ -312,7 +310,7 @@ def run_test(
 
         output_dtypes = None
         if check_dtype:
-            _, output_dtypes = infer_module_output_shapes_dtypes(
+            _, output_dtypes = infer_module_outputs(
                 mod,
                 input_specs,
                 compilation_settings.device,
@@ -407,7 +405,7 @@ def run_test_with_dynamic_shape(
         )
 
         if check_dtype:
-            _, output_dtypes = infer_module_output_shapes_dtypes(
+            _, output_dtypes = infer_module_outputs(
                 mod,
                 input_specs,
                 compilation_settings.device,