chore: cherry pick of #2832 (#2852)

zewenli98 · web-flow · commit 856f33d79ebd · 2024-05-29T14:27:56.000-07:00
diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py
@@ -317,8 +317,10 @@ def run(
         )
         timing_cache = self._create_timing_cache(builder_config, existing_cache)
 
-        engine = self.builder.build_serialized_network(self.ctx.net, builder_config)
-        assert engine
+        serialized_engine = self.builder.build_serialized_network(
+            self.ctx.net, builder_config
+        )
+        assert serialized_engine
 
         serialized_cache = (
             bytearray(timing_cache.serialize())
@@ -328,10 +330,10 @@ def run(
         _LOGGER.info(
             f"Build TRT engine elapsed time: {datetime.now() - build_engine_start_time}"
         )
-        _LOGGER.info(f"TRT Engine uses: {engine.nbytes} bytes of Memory")
+        _LOGGER.info(f"TRT Engine uses: {serialized_engine.nbytes} bytes of Memory")
 
         return TRTInterpreterResult(
-            engine, self._input_names, self._output_names, serialized_cache
+            serialized_engine, self._input_names, self._output_names, serialized_cache
         )
 
     def run_node(self, n: torch.fx.Node) -> torch.fx.Node:
diff --git a/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py
@@ -30,7 +30,7 @@ class PythonTorchTensorRTModule(Module):  # type: ignore[misc]
 
     def __init__(
         self,
-        engine: trt.ICudaEngine,
+        engine: bytes,
         input_names: Optional[List[str]] = None,
         output_names: Optional[List[str]] = None,
         target_device: Device = Device._current_device(),
@@ -61,9 +61,9 @@ def _initialize(self) -> None:
         self.engine = runtime.deserialize_cuda_engine(self.engine)
         self.context = self.engine.create_execution_context()
 
-        assert (
-            self.engine.num_io_tensors // self.engine.num_optimization_profiles
-        ) == (len(self.input_names) + len(self.output_names))
+        assert self.engine.num_io_tensors == (
+            len(self.input_names) + len(self.output_names)
+        )
 
         self.input_dtypes = [
             dtype._from(self.engine.get_tensor_dtype(input_name))
