content: import logging
  from typing import Any
  
  import torch
  from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
  
  logger = logging.getLogger(__name__)
  
  
  class _PreAllocatedOutputContextManager(object):
      """
      Helper class used to enable pre-allocated output feature in runtime module
      """
  
      def __init__(self, module: torch.fx.GraphModule) -> None:
          rt_mods = []
          for name, rt_mod in module.named_children():
              if "_run_on_acc" in name and isinstance(
                  rt_mod, (PythonTorchTensorRTModule, TorchTensorRTModule)
              ):
                  rt_mods.append(rt_mod)
          self.rt_mods = rt_mods
  
      def set_pre_allocated_output(self, enable: bool) -> None:
          for mod in self.rt_mods:
              mod.set_pre_allocated_outputs(enable)
  
      def __enter__(self) -> "_PreAllocatedOutputContextManager":
          # Enable pre-allocated output
          self.set_pre_allocated_output(True)
          return self
  
      def __exit__(self, *args: Any) -> None:
          # Disable pre-allocated output
          self.set_pre_allocated_output(False)
  
  
  def enable_pre_allocated_outputs(
      module: torch.fx.GraphModule,
  ) -> _PreAllocatedOutputContextManager:
      return _PreAllocatedOutputContextManager(module)
  