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 "_run_on_acc" not in str(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

setup.py

@@ -195,10 +195,10 @@ def build_libtorchtrt_cxx11_abi(
     else:
         cmd.append("//:libtorchtrt")
 
-    # if develop:
-    #     cmd.append("--compilation_mode=dbg")
-    # else:
-    cmd.append("--compilation_mode=opt")
+    if develop:
+        cmd.append("--compilation_mode=dbg")
+    else:
+        cmd.append("--compilation_mode=opt")
     if use_dist_dir:
         if IS_AARCH64:
             cmd.append("--distdir=third_party/dist_dir/aarch64-linux-gnu")