Skip to content

Commit 23f6058

Browse files
JacobSzwejbkafacebook-github-bot
authored andcommitted
save api (#7097)
Summary: Tired of writing out this python boilerplate Reviewed By: tarun292, lucylq Differential Revision: D66523473
1 parent a9565aa commit 23f6058

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
@@ -1519,3 +1519,17 @@ def write_to_file(self, open_file: io.BufferedIOBase) -> None:
15191519
reducing the peak memory usage.
15201520
"""
15211521
self._pte_data.write_to_file(open_file)
1522+
1523+
def save(self, path: str) -> None:
1524+
"""
1525+
Saves the serialized ExecuTorch binary to the file at `path`.
1526+
"""
1527+
if path[-4:] != ".pte":
1528+
logging.error(f"Path {path} does not end with .pte")
1529+
raise ValueError(f"Path {path} does not end with .pte")
1530+
try:
1531+
with open(path, "wb") as file:
1532+
self.write_to_file(file)
1533+
logging.info(f"Saved exported program to {path}")
1534+
except Exception as e:
1535+
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") # not .pte so should die

0 commit comments

Comments
 (0)