Fixed bugs

Author: cehongwang

py/torch_tensorrt/dynamo/_compiler.py

@@ -1071,8 +1071,6 @@ def preserve_module_specs(
                         f.write(trt_module.get_layer_info())
 
     # Only set the requires_unique_output flag for the last TRT Module when user has access to the output tensor
-    if trt_module:
-        trt_module.set_output_tensors_as_unowned(True)
 
     # Parse the graph I/O and store it in dryrun tracker
     parse_graph_io(gm, dryrun_tracker)
@@ -1081,7 +1079,11 @@ def preserve_module_specs(
     for name, trt_module in trt_modules.items():
         setattr(partitioned_module, name, trt_module)
         if settings.lazy_engine_init and not settings.enable_cross_compile_for_windows:
-            getattr(partitioned_module, name).setup_engine()
+            trt_module = getattr(partitioned_module, name)
+            trt_module.setup_engine()
+
+    if trt_module:
+        trt_module.set_output_tensors_as_unowned(True)
 
     # Reset settings object to user specification after fallback to global partitioning mode
     if fast_partitioner_failed:

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

@@ -225,10 +225,6 @@ def __init__(
         self.output_tensors_are_unowned = False
         if self.serialized_engine is not None and not self.settings.lazy_engine_init:
             self.setup_engine()
-        self.is_shape_inference_io = {
-            input_name: self.engine.is_shape_inference_io(input_name)
-            for input_name in self.input_names
-        }
 
     def set_output_tensors_as_unowned(self, enabled: bool) -> None:
         """
@@ -327,6 +323,11 @@ def setup_engine(self) -> None:
         if torch_tensorrt.runtime.get_cudagraphs_mode():
             self.cudagraph = torch.cuda.CUDAGraph()
 
+        self.is_shape_inference_io = {
+            input_name: self.engine.is_shape_inference_io(input_name)
+            for input_name in self.input_names
+        }
+
     def _check_initialized(self) -> None:
         if not self.initialized:
             raise RuntimeError("PythonTorchTensorRTModule is not initialized.")
@@ -517,7 +518,7 @@ def run_standard_execution() -> torch.Tensor | Tuple[torch.Tensor, ...]:
                 if can_use_pre_allocated_outputs:
                     outputs = self.pre_allocated_outputs
                 else:
-                    if shape_changed:
+                    if shape_changed or self.output_tensors is None:
                         self.output_shapes = [
                             tuple(self.context.get_tensor_shape(output_name))
                             for output_name in self.output_names