Skip to content

Commit 960b372

Browse files
committed
chore: Accept graphmodule as input to dynamo.compile
Signed-off-by: Dheeraj Peri <[email protected]>
1 parent 123a486 commit 960b372

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

py/torch_tensorrt/dynamo/compile.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,8 +85,10 @@ def compile(
8585
# Prepare torch_trt inputs
8686
inputs = prepare_inputs(inputs)
8787
device = to_torch_tensorrt_device(device)
88-
89-
gm = exported_program.module()
88+
if isinstance(exported_program, torch.fx.GraphModule):
89+
gm = exported_program
90+
else:
91+
gm = exported_program.module()
9092
logger.debug("Input graph: " + str(gm.graph))
9193

9294
# Apply lowering on the graph module

0 commit comments

Comments
 (0)