OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2026-01-15 12:15:41 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Merge pull request #49763 from geetachavan1/cherrypicks_50B03

Browse Source 
Prevent infinite loop/stack overflow in TFLite `while` op.
This commit is contained in:
Mihai Maruseac
2021-05-30 06:14:32 -07:00
committed by GitHub
parent a21b318bf5 7b3b84d8fc
commit 81d987f512
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/lite/kernels/while.cc
View File

@@ -132,6 +132,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();