Fixed the comments

Author: cehongwang

core/runtime/TRTEngine.cpp

@@ -289,12 +289,12 @@ void TRTEngine::enable_profiling() {
   exec_ctx->setProfiler(trt_engine_profiler.get());
 }
 
-void TRTEngine::set_requires_new_output_tensor(bool enable) {
-  this->requires_new_output_tensor = enable;
+void TRTEngine::set_unowned_output_tensor(bool enable) {
+  this->unowned_output_tensor = enable;
 }
 
-bool TRTEngine::get_requires_new_output_tensor() {
-  return this->requires_new_output_tensor;
+bool TRTEngine::is_unowned_output_tensor() {
+  return this->unowned_output_tensor;
 }
 
 void TRTEngine::set_profile_format(std::string format) {

core/runtime/TRTEngine.h

@@ -105,7 +105,7 @@ struct TRTEngine : torch::CustomClassHolder {
   std::pair<uint64_t, uint64_t> num_io;
   uint64_t io_size;
   std::map<std::string, bool> isShapeInferenceIO;
-  bool requires_new_output_tensor = false;
+  bool unowned_output_tensor = false;
   std::string name;
   RTDevice device_info;
 
@@ -162,8 +162,8 @@ struct TRTEngine : torch::CustomClassHolder {
   int64_t get_automatic_device_memory_budget();
   std::vector<at::Tensor> infer_outputs(std::vector<std::vector<int64_t>> input_shapes);
   void set_pre_allocated_outputs(bool enable);
-  void set_requires_new_output_tensor(bool enable);
-  bool get_requires_new_output_tensor();
+  void set_unowned_output_tensor(bool enable);
+  bool is_unowned_output_tensor();
   TorchTRTRuntimeStates runtime_states;
   friend std::ostream& operator<<(std::ostream& os, const TRTEngine& engine);
   static const char BINDING_DELIM = '%';

core/runtime/execute_engine.cpp

@@ -249,8 +249,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
       if (can_use_pre_allocated_outputs) {
         outputs = compiled_engine->pre_allocated_outputs;
       } else {
-        if (compiled_engine->allocated_outputs.size() == 0 or compiled_engine->requires_new_output_tensor or
-            shape_changed) {
+        if (compiled_engine->allocated_outputs.size() == 0 or compiled_engine->unowned_output_tensor or shape_changed) {
           compiled_engine->allocated_outputs = create_output_tensors(compiled_engine);
           new_outputs = true;
         }

core/runtime/register_jit_hooks.cpp

@@ -90,8 +90,8 @@ static auto TORCHTRT_UNUSED TRTEngineTSRegistrtion =
         .def("get_engine_layer_info", &TRTEngine::get_engine_layer_info)
         .def("infer_outputs", &TRTEngine::infer_outputs)
         .def("reset_captured_graph", &TRTEngine::reset_captured_graph)
-        .def("set_requires_new_output_tensor", &TRTEngine::set_requires_new_output_tensor)
-        .def("get_requires_new_output_tensor", &TRTEngine::get_requires_new_output_tensor)
+        .def("set_unowned_output_tensor", &TRTEngine::set_unowned_output_tensor)
+        .def("is_unowned_output_tensor", &TRTEngine::is_unowned_output_tensor)
         .def_readwrite("use_pre_allocated_outputs", &TRTEngine::use_pre_allocated_outputs)
         .def_readwrite("use_output_allocator_outputs", &TRTEngine::use_output_allocator_outputs)
         .def_property(

py/torch_tensorrt/dynamo/_compiler.py

@@ -996,7 +996,7 @@ def preserve_module_specs(
 
     # 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_requires_new_output_tensor(True)
+        trt_module.set_unowned_output_tensor(True)
 
     # Parse the graph I/O and store it in dryrun tracker
     parse_graph_io(gm, dryrun_tracker)

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

@@ -221,16 +221,28 @@ def __init__(
         self.use_output_allocator_outputs = False
         self.device = torch.cuda.current_device()
         self.cudagraphs_enabled = torch_tensorrt.runtime.get_cudagraphs_mode()
-        self.requires_new_output_tensor = False
+        # If the output tensor is not owned by the engine (unowned_output_tensor=True), we need to create a new output tensor in each forward pass
+        self.unowned_output_tensor = 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_requires_new_output_tensor(self, enabled: bool) -> None:
-        self.requires_new_output_tensor = enabled
+    def set_unowned_output_tensor(self, enabled: bool) -> None:
+        """
+        Set the flag to indicate if the output tensor is unowned by the engine.
+        If self.unowned_output_tensor=True, the engine will create a new output tensor in each forward pass.
+        This would be slower but is required when users need to manipulate the output tensor after each forward pass.
+        Therefore, this should be set to True only for the last module in a graph and leave to False for intermediate modules,
+        which users don't have access to.
+        Args:
+            enabled: bool
+                Whether to set the flag to True.
+
+        """
+        self.unowned_output_tensor = enabled
 
     def get_streamable_device_memory_budget(self) -> Any:
         return self.engine.streamable_weights_size
@@ -520,7 +532,7 @@ def run_standard_execution() -> torch.Tensor | Tuple[torch.Tensor, ...]:
                         )
                     if (
                         self.output_tensors is None
-                        or self.requires_new_output_tensor
+                        or self.unowned_output_tensor
                         or shape_changed
                     ):
                         self.output_tensors = self.create_output_tensors()

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -289,8 +289,8 @@ def set_extra_state(self, state: SerializedTorchTensorRTModuleFmt) -> None:
             metadata = TorchTensorRTModule.decode_metadata(serialized_metadata)
             self.settings = metadata["settings"]
             self.weight_name_map = metadata["weight_name_map"]
-            self.requires_new_output_tensor = metadata["requires_new_output_tensor"]
-            self.engine.set_requires_new_output_tensor(self.requires_new_output_tensor)
+            self.unowned_output_tensor = metadata["unowned_output_tensor"]
+            self.engine.set_unowned_output_tensor(self.unowned_output_tensor)
 
         else:
             self.engine = None
@@ -362,11 +362,11 @@ def enable_profiling(
         self.engine.enable_profiling()
         self.engine.set_profile_format(profile_format)
 
-    def set_requires_new_output_tensor(self, enabled: bool) -> None:
-        self.engine.set_requires_new_output_tensor(enabled)
+    def set_unowned_output_tensor(self, enabled: bool) -> None:
+        self.engine.set_unowned_output_tensor(enabled)
 
-    def get_requires_new_output_tensor(self) -> bool:
-        return self.engine.get_requires_new_output_tensor()
+    def is_unowned_output_tensor(self) -> bool:
+        return self.engine.is_unowned_output_tensor()
 
     def disable_profiling(self) -> None:
         """Disable the profiler"""