27 lines
1.1 KiB
Diff
27 lines
1.1 KiB
Diff
|
From a74768f8e4efbda4def9f16ee7e13cf3922ac5f7 Mon Sep 17 00:00:00 2001
|
||
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
||
|
|
Date: Thu, 6 May 2021 14:24:09 -0700
|
||
|
|
Subject: [PATCH] Prevent heap OOB error in `MaxPoolGrad`
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 372424854
|
||
|
|
Change-Id: Idac0f23867ad8b0601cafbaaa52d5e64269e63a7
|
||
|
|
---
|
||
|
|
tensorflow/core/kernels/maxpooling_op.cc | 4 +++-
|
||
|
|
1 file changed, 3 insertions(+), 1 deletion(-)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/maxpooling_op.cc b/tensorflow/core/kernels/maxpooling_op.cc
|
||
|
|
index ceb6694ed665d..01f303eb26980 100644
|
||
|
|
--- a/tensorflow/core/kernels/maxpooling_op.cc
|
||
|
|
+++ b/tensorflow/core/kernels/maxpooling_op.cc
|
||
|
|
@@ -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);
|
||
|
|
+ }
|
||
|
|
}
|
||
|
|
}
|
||
|
|
};
|