fixed redundant code in TRT Interpreter

pytorch · Feb 14, 2025 · 6824e57 · 6824e57
1 parent 2368e63
commit 6824e57
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 8 deletions.
diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py
@@ -33,7 +33,9 @@
 from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
     DYNAMO_CONVERTERS as CONVERTERS,
 )
-from torch_tensorrt.dynamo.conversion._ConverterRegistry import CallingConvention
+from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
+    CallingConvention,
+)
 from torch_tensorrt.dynamo.conversion._TRTBuilderMonitor import TRTBulderMonitor
 from torch_tensorrt.dynamo.conversion.converter_utils import (
     get_node_io,
@@ -740,10 +742,6 @@ def run(
     def run_node(self, n: torch.fx.Node) -> torch.fx.Node:
         self._cur_node_name = get_node_name(n)
         self._cur_node = n
-        # add "_itensor_to_tensor_meta"
-        kwargs = dict(n.kwargs)
-        kwargs["_itensor_to_tensor_meta"] = self._itensor_to_tensor_meta
-        n.kwargs = kwargs
 
         if _LOGGER.isEnabledFor(logging.DEBUG):
             _LOGGER.debug(
@@ -759,9 +757,7 @@ def run_node(self, n: torch.fx.Node) -> torch.fx.Node:
             f"Converted node {self._cur_node_name} [{n.target}] ({get_node_io(n, self.const_mapping)})"
         )
 
-        # remove "_itensor_to_tensor_meta"
         kwargs = dict(n.kwargs)
-        del kwargs["_itensor_to_tensor_meta"]
         n.kwargs = kwargs
 
         if isinstance(trt_node, trt.ITensor):

diff --git a/py/torch_tensorrt/fx/converters/acc_ops_converters.py b/py/torch_tensorrt/fx/converters/acc_ops_converters.py
@@ -3221,7 +3221,7 @@ def acc_ops_dequantize(
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     input_val = kwargs["input"]
-    input_val_tensor_meta = kwargs["_itensor_to_tensor_meta"][input_val]  # type: ignore[index]
+    input_val_tensor_meta = network._itensor_to_tensor_meta[input_val]  # type: ignore[index]
 
     if not isinstance(input_val, TRTTensor):
         raise RuntimeError(