From b9d6a4c752563a6863fa99f6c39922519e297260 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Thu, 14 Aug 2025 01:33:05 +0800 Subject: [PATCH] fix: fix export in torchscript format. --- py/torch_tensorrt/_compile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index acae618f1b..b295dbe34d 100644 --- 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" )