File tree Expand file tree Collapse file tree 1 file changed +7
-7
lines changed Expand file tree Collapse file tree 1 file changed +7
-7
lines changed Original file line number Diff line number Diff line change @@ -454,17 +454,17 @@ def _get_new_signature( # noqa: C901
454
454
new_state_dict = {}
455
455
new_constants = {}
456
456
457
- input_tensor_node_to_sig = {
458
- input_spec . arg . name : input_spec
459
- for input_spec in old_signature . input_specs
460
- if isinstance ( input_spec . arg , TensorArgument )
461
- }
457
+ placeholder_nodes = [
458
+ node . name for node in original_program . graph . nodes if node . op == "placeholder"
459
+ ]
460
+ assert len ( placeholder_nodes ) == len ( old_signature . input_specs )
461
+ input_node_to_sig = dict ( zip ( placeholder_nodes , old_signature . input_specs ))
462
462
463
463
for node in gm .graph .nodes :
464
464
is_tagged = tag is None or node .meta .get ("delegation_tag" , None ) == tag
465
465
if node .op == "placeholder" :
466
466
467
- if node .name not in input_tensor_node_to_sig :
467
+ if node .name not in input_node_to_sig :
468
468
assert tag is not None
469
469
input_specs .append (
470
470
InputSpec (
@@ -475,7 +475,7 @@ def _get_new_signature( # noqa: C901
475
475
)
476
476
continue
477
477
478
- orig_input_spec = input_tensor_node_to_sig [node .name ]
478
+ orig_input_spec = input_node_to_sig [node .name ]
479
479
480
480
if not isinstance (orig_input_spec .arg , TensorArgument ):
481
481
input_specs .append (orig_input_spec )
You can’t perform that action at this time.
0 commit comments