Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent heap OOB error in MaxPoolGrad
PiperOrigin-RevId: 372424854
Change-Id: Idac0f23867ad8b0601cafbaaa52d5e64269e63a7
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed May 6, 2021
1 parent ab1e644 commit a74768f
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion tensorflow/core/kernels/maxpooling_op.cc
Expand Up @@ -199,7 +199,9 @@ static void SpatialMaxPoolWithArgMaxHelper(
// CHECK(input_backprop_index >= in_start && input_backprop_index <
// in_end)
FastBoundsCheck(input_backprop_index - in_start, in_end - in_start);
input_backprop_flat(input_backprop_index) += out_backprop_flat(index);
if (index < out_backprop.NumElements()) {
input_backprop_flat(input_backprop_index) += out_backprop_flat(index);
}
}
}
};
Expand Down

0 comments on commit a74768f

Please sign in to comment.