From 5fb4787d966001e24f9e878fb26a75b60d23885c Mon Sep 17 00:00:00 2001 From: "Yanan Cao (PyTorch)" Date: Thu, 19 Dec 2024 10:56:56 -0800 Subject: [PATCH] executorch/test/models (#7387) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/7387 Reviewed By: avikchaudhuri, ydwu4 Differential Revision: D67384250 --- test/models/export_delegated_program.py | 6 ++++-- test/models/generate_linear_out_bundled_program.py | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/test/models/export_delegated_program.py b/test/models/export_delegated_program.py index a37fe32e556..a85dab6753f 100644 --- a/test/models/export_delegated_program.py +++ b/test/models/export_delegated_program.py @@ -130,7 +130,9 @@ def __init__(self, fn): def forward(self, *args, **kwargs): return self.fn(*args, **kwargs) - exported_program = export(WrapperModule(getattr(eager_module, method)), args=inputs) + exported_program = export( + WrapperModule(getattr(eager_module, method)), args=inputs, strict=True + ) edge_config = EdgeCompileConfig(_check_ir_validity=False) et_config = exir.ExecutorchBackendConfig( @@ -167,7 +169,7 @@ def forward(self, *args, **kwargs): composite_module(*inputs) executorch_program = to_edge( - export(composite_module, args=inputs) + export(composite_module, args=inputs, strict=True) ).to_executorch(config=et_config) return executorch_program.buffer diff --git a/test/models/generate_linear_out_bundled_program.py b/test/models/generate_linear_out_bundled_program.py index 93fd1445ef5..c98ea7ed683 100644 --- a/test/models/generate_linear_out_bundled_program.py +++ b/test/models/generate_linear_out_bundled_program.py @@ -37,7 +37,7 @@ def main() -> None: trace_inputs = (torch.ones(2, 2, dtype=torch.float),) # Trace to FX Graph. - exec_prog = to_edge(export(model, trace_inputs)).to_executorch( + exec_prog = to_edge(export(model, trace_inputs, strict=True)).to_executorch( config=ExecutorchBackendConfig( memory_planning_pass=MemoryPlanningPass(), to_out_var_pass=ToOutVarPass(),