Skip to content

Commit 3fcf0bd

Browse files
save api
Differential Revision: D66523473 Pull Request resolved: #7097
1 parent 8460d42 commit 3fcf0bd

File tree

2 files changed

+23
-1
lines changed

2 files changed

+23
-1
lines changed

exir/program/_program.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1532,3 +1532,17 @@ def write_to_file(self, open_file: io.BufferedIOBase) -> None:
15321532
reducing the peak memory usage.
15331533
"""
15341534
self._pte_data.write_to_file(open_file)
1535+
1536+
def save(self, path: str) -> None:
1537+
"""
1538+
Saves the serialized ExecuTorch binary to the file at `path`.
1539+
"""
1540+
if path[-4:] != ".pte":
1541+
logging.error(f"Path {path} does not end with .pte")
1542+
raise ValueError(f"Path {path} does not end with .pte")
1543+
try:
1544+
with open(path, "wb") as file:
1545+
self.write_to_file(file)
1546+
logging.info(f"Saved exported program to {path}")
1547+
except Exception as e:
1548+
logging.error(f"Error while saving to {path}: {e}")

exir/program/test/test_program.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7-
# pye-strict
7+
# pyre-unsafe
88

99
import copy
1010
import unittest
@@ -803,3 +803,11 @@ def test_to_edge_with_preserved_ops_not_in_model(self):
803803
self._test_to_edge_with_preserved_ops(
804804
program, ops_not_to_decompose, expected_non_decomposed_edge_ops
805805
)
806+
807+
def test_save_fails(self):
808+
model = TestLinear()
809+
program = torch.export.export(model, model._get_random_inputs())
810+
edge = to_edge(program)
811+
et = edge.to_executorch()
812+
with self.assertRaises(ValueError):
813+
_ = et.save("/tmp/test_save.pt")

0 commit comments

Comments
 (0)