fix: fix export in torchscript format.

GdoongMathew · web-flow · commit b9d6a4c75256 · 2025-08-14T01:33:05.000+08:00
diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py
@@ -639,7 +639,7 @@ def save(
             "Input model is of type nn.Module. Saving nn.Module directly is not supported. Supported model types torch.jit.ScriptModule | torch.fx.GraphModule | torch.export.ExportedProgram."
         )
     elif module_type == _ModuleType.ts:
-        if not all([output_format == f for f in ["exported_program", "aot_inductor"]]):
+        if output_format != "torchscript":
             raise ValueError(
                 "Provided model is a torch.jit.ScriptModule but the output_format specified is not torchscript. Other output formats are not supported"
             )

Original file line number	Diff line number	Diff line change
`@@ -639,7 +639,7 @@ def save(`
`639`	`639`	`"Input model is of type nn.Module. Saving nn.Module directly is not supported. Supported model types torch.jit.ScriptModule \| torch.fx.GraphModule \| torch.export.ExportedProgram."`
`640`	`640`	`)`
`641`	`641`	`elif module_type == _ModuleType.ts:`
`642`		`- if not all([output_format == f for f in ["exported_program", "aot_inductor"]]):`
	`642`	`+ if output_format != "torchscript":`
`643`	`643`	`raise ValueError(`
`644`	`644`	`"Provided model is a torch.jit.ScriptModule but the output_format specified is not torchscript. Other output formats are not supported"`
`645`	`645`	`)`