Skip to content

Commit d55c451

Browse files
committed
disable cuda graph for output allocator mode
1 parent 8ef1a87 commit d55c451

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

+1
Original file line numberDiff line numberDiff line change
@@ -545,6 +545,7 @@ def run_cuda_graph() -> torch.Tensor | Tuple[torch.Tensor, ...]:
545545
return outputs
546546

547547
def run_output_allocator() -> torch.Tensor | Tuple[torch.Tensor, ...]:
548+
torch_tensorrt.runtime.set_cudagraphs_mode(False)
548549
with (
549550
torch.autograd.profiler.record_function(
550551
"PythonTorchTensorRTModule:ProcessInputs"

0 commit comments

Comments
 (0)