Fix test

jackzhxng · jackzhxng · commit a6a716239d40 · 2024-11-06T20:02:14.000-05:00
diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py
@@ -203,7 +203,7 @@ def export(self) -> "LLMEdgeManager":
                     dynamic_shapes=dynamic_shape,
                 )
             self.pre_autograd_graph_module = exported_module.module()
-            if self.args.export_only:
+            if hasattr(self.args, "export_only") and self.args.export_only:
                 torch.export.save(exported_module, self.args.output_name)
 
         return self

Original file line number	Diff line number	Diff line change
`@@ -203,7 +203,7 @@ def export(self) -> "LLMEdgeManager":`
`203`	`203`	`dynamic_shapes=dynamic_shape,`
`204`	`204`	`)`
`205`	`205`	`self.pre_autograd_graph_module = exported_module.module()`
`206`		`- if self.args.export_only:`
	`206`	`+ if hasattr(self.args, "export_only") and self.args.export_only:`
`207`	`207`	`torch.export.save(exported_module, self.args.output_name)`
`208`	`208`
`209`	`209`	`return self`