Skip to content

Commit 19818af

Browse files
v-dziubacopybara-github
authored andcommitted
Avoid adding input/output operators if they are already in the subgraph.
PiperOrigin-RevId: 764813584
1 parent 1b04849 commit 19818af

File tree

1 file changed

+9
-4
lines changed

1 file changed

+9
-4
lines changed

ai_edge_quantizer/calibrator.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -140,10 +140,15 @@ def calibrate(
140140
graph_info = qtyping.GraphInfo(
141141
subgraph.tensors, self._flatbuffer_model.buffers
142142
)
143-
# Add input/output operators to the subgraph.
144-
subgraph.operators += (
145-
tfl_flatbuffer_utils.get_subgraph_input_output_operators(subgraph)
146-
)
143+
# Add input/output operators if they are not in the subgraph.
144+
if not any(
145+
isinstance(op, qtyping.IOOperator) for op in subgraph.operators
146+
):
147+
subgraph.operators += (
148+
tfl_flatbuffer_utils.get_subgraph_input_output_operators(
149+
subgraph
150+
)
151+
)
147152
for op in subgraph.operators:
148153
if isinstance(op, qtyping.IOOperator):
149154
op_key = op.op_key

0 commit comments

Comments
 (0)