diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index 7683d9c44d1..b8223d84f51 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -117,7 +117,7 @@ def forward(self, q, k, v, mask): v = torch.randn(batch_size, n_heads, max_seq_length, embedding_dim) mask = torch.randn(seq_len, max_seq_length) example_inputs = (q, k, v, mask) - ep = torch.export.export(model, example_inputs) + ep = torch.export.export(model, example_inputs, strict=True) coreml_partitioner = CoreMLPartitioner() # Using to_edge_transform_and_lower, we expect SDPA will be preserved and show up in delegated graph diff --git a/backends/qualcomm/_passes/decompose_any.py b/backends/qualcomm/_passes/decompose_any.py index c0c65ee7040..e92bf11dd18 100644 --- a/backends/qualcomm/_passes/decompose_any.py +++ b/backends/qualcomm/_passes/decompose_any.py @@ -41,7 +41,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: keepdim = node.args[2] if len(node.args) > 2 else False model = Any(dim, keepdim) edge_mgr = to_edge( - torch.export.export(model, (node.args[0].meta["val"],)) + torch.export.export(model, (node.args[0].meta["val"],), strict=True) ) decomposed_module = edge_mgr.exported_program() diff --git a/backends/qualcomm/_passes/decompose_linalg_vector_norm.py b/backends/qualcomm/_passes/decompose_linalg_vector_norm.py index 0ee74720c78..4a54c2aa50c 100644 --- a/backends/qualcomm/_passes/decompose_linalg_vector_norm.py +++ b/backends/qualcomm/_passes/decompose_linalg_vector_norm.py @@ -46,11 +46,13 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: model = LinalgVectorNorm(ord, dim, keepdim) if self.aten_dialect_capture: decomposed_module = torch.export.export( - model, (node.args[0].meta["val"],) + model, (node.args[0].meta["val"],), strict=True ).module() else: edge_mgr = to_edge( - torch.export.export(model, (node.args[0].meta["val"],)) + torch.export.export( + model, (node.args[0].meta["val"],), strict=True + ) ) decomposed_module = edge_mgr.exported_program() diff --git a/backends/qualcomm/tests/utils.py b/backends/qualcomm/tests/utils.py index 4908bf889a9..100469a72c3 100644 --- a/backends/qualcomm/tests/utils.py +++ b/backends/qualcomm/tests/utils.py @@ -526,7 +526,9 @@ def get_qdq_module( dynamic_shapes: Dict = None, bypass_check: bool = False, ) -> torch.fx.GraphModule: - m = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes).module() + m = torch.export.export( + module, inputs, dynamic_shapes=dynamic_shapes, strict=True + ).module() quantizer = QnnQuantizer() quantizer.add_custom_quant_annotations(custom_quant_annotations) diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index 5ae640adc6e..edd674d6381 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -452,7 +452,7 @@ def capture_program( dynamic_shapes: Dict = None, ) -> exir.ExirExportedProgram: module = _preprocess_module(module, inputs) - ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes) + ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes, strict=True) decomposed_ep = ep.run_decompositions(get_decomp_table()) core_ep = ExirExportedProgram(decomposed_ep, False) core_ep.transform(TensorI64toI32(edge_program=core_ep)) diff --git a/backends/transforms/test/test_rank_0_to_rank_1.py b/backends/transforms/test/test_rank_0_to_rank_1.py index 50c6357fb67..eddad536e06 100644 --- a/backends/transforms/test/test_rank_0_to_rank_1.py +++ b/backends/transforms/test/test_rank_0_to_rank_1.py @@ -17,7 +17,7 @@ def forward(self, x, y): model.eval() example_inputs = (torch.tensor(1.0), torch.tensor(2.0)) - aten = torch.export.export(model, example_inputs) + aten = torch.export.export(model, example_inputs, strict=True) # Check that the input rank is 0 for node in aten.graph.nodes: diff --git a/examples/apple/coreml/llama/export.py b/examples/apple/coreml/llama/export.py index b39162c484b..9aa232fa691 100644 --- a/examples/apple/coreml/llama/export.py +++ b/examples/apple/coreml/llama/export.py @@ -192,10 +192,7 @@ def main() -> None: ) example_inputs = input_manager.get_inputs(tokens=[0]) - ep = torch.export.export( - model, - example_inputs, - ) + ep = torch.export.export(model, example_inputs, strict=True) print("Exported program") print(ep) diff --git a/exir/tests/test_memory_planning.py b/exir/tests/test_memory_planning.py index d5e5627dfa6..f3cddaff643 100644 --- a/exir/tests/test_memory_planning.py +++ b/exir/tests/test_memory_planning.py @@ -749,7 +749,7 @@ def forward(self, input, label): net = TrainingNet(Net()) inputs = (torch.randn(1, 6, 5, 5), torch.ones(1, dtype=torch.int64)) - ep = export(net, inputs) + ep = export(net, inputs, strict=True) ep = _export_forward_backward(ep) ep = to_edge(ep) ep = ep.to_executorch() diff --git a/exir/verification/test/test_verifier.py b/exir/verification/test/test_verifier.py index b2182242dd7..8520d3ce13e 100644 --- a/exir/verification/test/test_verifier.py +++ b/exir/verification/test/test_verifier.py @@ -153,7 +153,7 @@ def forward(self, input, label): net = TrainingNet(Net()) inputs = (torch.randn(1, 6, 5, 5), torch.ones(1, dtype=torch.int64)) - export_model = export(net, inputs) + export_model = export(net, inputs, strict=True) export_model = _export_forward_backward(export_model) edge = to_edge(export_model)