Skip to content

Commit ee1dc06

Browse files
committed
Enable CG tests Python + Fallback
1 parent d45c23c commit ee1dc06

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

tests/py/dynamo/runtime/test_cudagraphs.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,10 @@ def forward(self, x):
134134
torch_executed_ops={"torch.ops.aten.mul.Tensor"},
135135
use_python_runtime=True,
136136
)
137-
optimized_model_results = optimized_model(*inputs).detach().cpu()
137+
138+
with torch_tensorrt.runtime.enable_cudagraphs():
139+
optimized_model_results = optimized_model(*inputs).detach().cpu()
140+
138141
torch_model_results = fx_graph(*inputs).detach().cpu()
139142

140143
max_diff = float(

0 commit comments

Comments
 (0)