@@ -178,18 +178,14 @@ def export_program(
178
178
return expo_program
179
179
180
180
181
- # Export the model and lower it to an EdgeProgramManager (in edge IR).
182
- def export_to_edge (
183
- model : torch .nn .Module ,
184
- inputs : tuple [object , ...],
181
+ def lower_ep_to_edge (
182
+ expo_program : ExportedProgram ,
185
183
dump_graphs : bool = False ,
186
184
constant_methods : Optional [dict [str , object ]] = None ,
187
185
) -> EdgeProgramManager :
188
- assert isinstance (model , torch .nn .Module ), "model should be an nn.Module"
189
-
190
- # Export the model into an ExportedProgram.
191
- expo_program = export_program (model , inputs )
192
-
186
+ """
187
+ Lower an ExportedProgram to an EdgeProgramManager (in edge IR).
188
+ """
193
189
# Call to_edge to convert the graph to edge IR.
194
190
# Note: dim_order is skipped (https://github.com/pytorch/executorch/issues/3704)
195
191
edge_prog_manager = to_edge (
@@ -215,6 +211,23 @@ def export_to_edge(
215
211
logging .info (
216
212
edge_prog_manager .exported_program ().graph_module .graph .print_tabular ()
217
213
)
214
+ return edge_prog_manager
215
+
216
+
217
+ # Export the model and lower it to an EdgeProgramManager (in edge IR).
218
+ def export_to_edge (
219
+ model : torch .nn .Module ,
220
+ inputs : tuple [object , ...],
221
+ dump_graphs : bool = False ,
222
+ constant_methods : Optional [dict [str , object ]] = None ,
223
+ ) -> EdgeProgramManager :
224
+ assert isinstance (model , torch .nn .Module ), "model should be an nn.Module"
225
+
226
+ # Export the model into an ExportedProgram.
227
+ expo_program = export_program (model , inputs )
228
+
229
+ # Lower the model to edge IR.
230
+ edge_prog_manager = lower_ep_to_edge (expo_program , dump_graphs , constant_methods )
218
231
219
232
return edge_prog_manager
220
233
0 commit comments