Changed the how unown_tensor attribute is set on TRT mod

Author: cehongwang

py/torch_tensorrt/dynamo/_compiler.py

@@ -949,7 +949,7 @@ def preserve_module_specs(
     for attr in dir(gm):
         if attr.startswith("_frozen_param"):
             delattr(gm, attr)
-    trt_module = None
+
     for name, _ in partitioned_module.named_children():
         submodule = getattr(partitioned_module, name)
         # filter on the GraphModule
@@ -1082,8 +1082,12 @@ def preserve_module_specs(
             trt_module = getattr(partitioned_module, name)
             trt_module.setup_engine()
 
-    if trt_module:
-        trt_module.set_output_tensors_as_unowned(True)
+    output_node = list(partitioned_module.graph.nodes)[-1]
+    for arg in output_node.args:
+        target = arg[0].target
+        if "acc" not in target:
+            continue
+        getattr(partitioned_module, target).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

@@ -807,3 +807,6 @@ def validate_input_shapes(self, inputs: Sequence[torch.Tensor]) -> bool:
             return True
 
         return False
+
+    def are_output_tensors_unowned(self) -> bool:
+        return self.output_tensors_are_unowned