Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix heap buffer overflow
PiperOrigin-RevId: 372132844
Change-Id: Idef9895efaf145f2b1c23d31983601ec980cd5e4
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed May 5, 2021
1 parent ef0c008 commit dcd7867
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions tensorflow/core/kernels/maxpooling_op.cc
Expand Up @@ -1014,6 +1014,9 @@ struct LaunchMaxPoolingGradWithArgmax<CPUDevice, T> {
const int input_start = start * input_size_per_batch;
const int input_end = limit * input_size_per_batch;
for (int64 index = input_start; index < input_end; index++) {
if (index >= argmax.NumElements()) {
break;
}
int64 grad_out_index = argmax_flat(index);
if (!include_batch_in_index) {
const int64 cur_batch = index / input_size_per_batch;
Expand Down

0 comments on commit dcd7867

Please sign in to comment.