additional operator relocation check.

rkuester · Oct 30, 2024 · fc4b473 · fc4b473
1 parent 8c53ee3
commit fc4b473
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/lite/micro/compression/relocate_ops.py b/tensorflow/lite/micro/compression/relocate_ops.py
@@ -145,6 +145,9 @@ def process_operator_var_handle(context: Context) -> VarHandles:
 
 def process_operator_assign_variable(context: Context) -> VarHandles:
   assign_op = context.current_op()
+  pending_concat_op = context.get_concat_op_by_tensor(
+      assign_op.inputs_indices[1], assign_op.subgraph.index)
+  assert pending_concat_op is None
   read_var_op = context.get_read_var_op_by_tensor(assign_op.inputs_indices[1],
                                                   assign_op.subgraph.index)
   if read_var_op is not None: