fix bugs

zewenli98 · zewenli98 · commit fa0162827b4f · 2024-05-13T17:37:07.000-07:00
diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py
@@ -313,8 +313,8 @@ 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())
@@ -324,10 +324,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
@@ -29,7 +29,7 @@ class PythonTorchTensorRTModule(Module):  # type: ignore[misc]
 
     def __init__(
         self,
-        engine: trt.ICudaEngine,
+        engine: trt.tensorrt.IHostMemory,
         input_names: Optional[List[str]] = None,
         output_names: Optional[List[str]] = None,
         target_device: Device = Device._current_device(),
@@ -60,9 +60,7 @@ 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))

Original file line number	Diff line number	Diff line change
`@@ -313,8 +313,8 @@ def run(`
`313`	`313`	`)`
`314`	`314`	`timing_cache = self._create_timing_cache(builder_config, existing_cache)`
`315`	`315`
`316`		`- engine = self.builder.build_serialized_network(self.ctx.net, builder_config)`
`317`		`- assert engine`
	`316`	`+ serialized_engine = self.builder.build_serialized_network(self.ctx.net, builder_config)`
	`317`	`+ assert serialized_engine`
`318`	`318`
`319`	`319`	`serialized_cache = (`
`320`	`320`	`bytearray(timing_cache.serialize())`
`@@ -324,10 +324,10 @@ def run(`
`324`	`324`	`_LOGGER.info(`
`325`	`325`	`f"Build TRT engine elapsed time: {datetime.now() - build_engine_start_time}"`
`326`	`326`	`)`
`327`		`- _LOGGER.info(f"TRT Engine uses: {engine.nbytes} bytes of Memory")`
	`327`	`+ _LOGGER.info(f"TRT Engine uses: {serialized_engine.nbytes} bytes of Memory")`
`328`	`328`
`329`	`329`	`return TRTInterpreterResult(`
`330`		`- engine, self._input_names, self._output_names, serialized_cache`
	`330`	`+ serialized_engine, self._input_names, self._output_names, serialized_cache`
`331`	`331`	`)`
`332`	`332`
`333`	`333`	`def run_node(self, n: torch.fx.Node) -> torch.fx.Node:`