disable cuda graph for output allocator mode

pytorch · Feb 14, 2025 · d55c451 · d55c451
1 parent 8ef1a87
commit d55c451
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py
@@ -545,6 +545,7 @@ def run_cuda_graph() -> torch.Tensor | Tuple[torch.Tensor, ...]:
             return outputs
 
         def run_output_allocator() -> torch.Tensor | Tuple[torch.Tensor, ...]:
+            torch_tensorrt.runtime.set_cudagraphs_mode(False)
             with (
                 torch.autograd.profiler.record_function(
                     "PythonTorchTensorRTModule:ProcessInputs"