File tree Expand file tree Collapse file tree 1 file changed +11
-2
lines changed Expand file tree Collapse file tree 1 file changed +11
-2
lines changed Original file line number Diff line number Diff line change @@ -261,6 +261,15 @@ def _partition_and_lower_one_graph_module(
261
261
call_delegate_args .append (inp_node )
262
262
break
263
263
264
+ def generate_debug_handle (ep : ExportedProgram ) -> int :
265
+ """
266
+ Generate a debug handle for the given ExportedProgram.
267
+ """
268
+ debug_handle = 0
269
+ for node in ep .graph_module .graph .nodes :
270
+ debug_handle = max (debug_handle , node .meta .get ("debug_handle" , 0 ))
271
+ return debug_handle + 1
272
+
264
273
# Replace the partitioned submodule with a lowered submodule
265
274
# Add call_method node with function "forward"
266
275
with tagged_graph_module .graph .inserting_before (call_module_node ):
@@ -273,8 +282,8 @@ def _partition_and_lower_one_graph_module(
273
282
(lowered_node ,) + tuple (call_delegate_args ),
274
283
call_module_node .kwargs ,
275
284
)
276
- call_delegate_node .meta ["debug_handle" ] = len (
277
- tagged_graph_module . graph . nodes
285
+ call_delegate_node .meta ["debug_handle" ] = generate_debug_handle (
286
+ owning_program
278
287
)
279
288
call_delegate_node .meta ["val" ] = submodule_output_node .meta ["val" ]
280
289
call_module_node .replace_all_uses_with (call_delegate_node )
You can’t perform that action at this time.
0 commit comments