Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent infinite loop/stack overflow in TFLite while op.
PiperOrigin-RevId: 370800333
Change-Id: I6a2e4ff849da339545c449db2af7e11ce6ff02c3
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Apr 28, 2021
1 parent 106d8f4 commit 9c1dc92
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions tensorflow/lite/kernels/while.cc
Expand Up @@ -138,6 +138,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
auto* subgraphs = this_subgraph->GetSubgraphs();
TF_LITE_ENSURE(context, op_data->cond_subgraph_index < subgraphs->size());
TF_LITE_ENSURE(context, op_data->body_subgraph_index < subgraphs->size());
TF_LITE_ENSURE(context,
op_data->cond_subgraph_index != op_data->body_subgraph_index);

Subgraph* cond_subgraph = (*subgraphs)[op_data->cond_subgraph_index].get();
Subgraph* body_subgraph = (*subgraphs)[op_data->body_subgraph_index].get();
Expand Down

0 comments on commit 9c1dc92

Please sign in to comment.