Fixed a small typo in the test file

Author: cehongwang

py/torch_tensorrt/dynamo/_compiler.py

@@ -622,9 +622,9 @@ def convert_module_to_trt_engine(
     if kwarg_inputs is None:
         kwarg_inputs = {}
     # Prepare torch_trt inputs
-    arg_inputs = prepare_inputs(arg_inputs)
+    arg_input_list = list(prepare_inputs(arg_inputs))
     kwarg_input_list = prepare_inputs(kwarg_inputs)
-    flattened_input_list = arg_inputs + flatten_dict_value(kwarg_input_list)
+    flattened_input_list = arg_input_list + flatten_dict_value(kwarg_input_list)
     device = to_torch_tensorrt_device(device)
     enabled_precisions = {dtype._from(e) for e in enabled_precisions}
 
@@ -678,7 +678,7 @@ def convert_module_to_trt_engine(
         interpreter_result = interpret_module_to_result(
             gm,
             inputs=flattened_input_list,
-            arg_inputs=arg_inputs,
+            arg_inputs=arg_input_list,
             kwarg_inputs=kwarg_input_list,
             settings=settings,
         )

py/torch_tensorrt/dynamo/_refit.py

@@ -157,7 +157,8 @@ def refit_module_weights(
                         This compiled_module should be compmiled by torch_tensorrt.dynamo.compile
                         or load it from disk using trt.load.
         new_weight_module: exported program with the updated weights. This one should have the same model architecture as the compiled module.
-        inputs: sample inputs
+        arg_inputs: sample arg inputs. Optional, needed if output check
+        kwarg_inputs: sample kwarg inputs. Optional, needed if output check
         verify_output: whether to verify output of refitted module
     Returns:
         A new compiled TensorRT module that has the updated weights.

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -4,6 +4,7 @@
 import logging
 from typing import Any, List, Optional, Sequence
 
+import tensorrt as trt
 import torch
 from torch.fx.experimental.proxy_tensor import maybe_disable_fake_tensor_mode
 from torch_tensorrt._Device import Device
@@ -18,8 +19,6 @@
 from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
 from torch_tensorrt.dynamo.utils import get_torch_inputs
 
-import tensorrt as trt
-
 logger = logging.getLogger(__name__)
 
 
@@ -30,6 +29,10 @@ def infer_module_output_dtypes(
     kwarg_inputs: Optional[dict[str, Any]] = None,
     truncate_double: bool = False,
 ) -> List[dtype]:
+    """
+    inputs can be either arg_inputs or flattened input list. If it is flattened list, kwarg_inputs
+    should be None, as it is already included in the flattened input.
+    """
     with maybe_disable_fake_tensor_mode():
         torch_inputs = get_torch_inputs(inputs, device)
         if kwarg_inputs is None:
@@ -72,7 +75,10 @@ def interpret_module_to_result(
     """Interpret an FX module to a TRTInterpreterResult
     Args:
         module: FX GraphModule to interpret
-        inputs: Sequence of Tensors representing inputs to the module
+        inputs: Sequence of FLATTENED Tensors representing inputs to the module. It should include both
+                arg_inputs and kwarg_inputs, if applicable.
+        arg_inputs: Sequence of Tensors representing inputs to the module.
+        kwarg_inputs: A dictionary of Tensors representing inputs to the module.
         settings: Compilation settings
     Returns:
         TRTInterpreterResult

tests/py/dynamo/models/test_export_kwargs_serde.py

@@ -453,7 +453,7 @@ def forward(self, x, b=None, c=None, d=None, e=[]):
         msg=f"CustomKwargs Module TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
     )
     # Change the input shape
-    kwargs["d"][1] = torch.randn([10, 2])
+    kwargs["e"][1] = torch.randn([10, 2]).to("cuda")
     cos_sim = cosine_similarity(model(*args, **kwargs), trt_gm(*args, **kwargs)[0])
     assertions.assertTrue(
         cos_sim > COSINE_THRESHOLD,