diff --git a/codegen/runtime/micro_codegen_context.cc b/codegen/runtime/micro_codegen_context.cc index 3d2e479be1f..858c823c1f3 100644 --- a/codegen/runtime/micro_codegen_context.cc +++ b/codegen/runtime/micro_codegen_context.cc @@ -95,6 +95,7 @@ TfLiteStatus MicroCodegenContext::InvokeSubgraph(int subgraph_idx) { TF_LITE_ENSURE(context_, static_cast(subgraph_idx) < subgraphs_.size()); size_t previous_subgraph_idx = current_subgraph_idx_; + current_subgraph_idx_ = subgraph_idx; TfLiteStatus status = subgraphs_[subgraph_idx].invoke(context_, subgraphs_[subgraph_idx].nodes); current_subgraph_idx_ = previous_subgraph_idx;