25 lines
1.0 KiB
Diff
25 lines
1.0 KiB
Diff
From 376c352a37ce5a68b721406dc7e77ac4b6cf483d Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
Date: Wed, 5 May 2021 14:34:54 -0700
|
|
Subject: [PATCH] Don't do any work if output tensor is null (prevent div by 0)
|
|
|
|
PiperOrigin-RevId: 372208700
|
|
Change-Id: Iea6b6293e887ade8538facfdb50fb931e17f511e
|
|
---
|
|
tensorflow/core/kernels/maxpooling_op.cc | 2 ++
|
|
1 file changed, 2 insertions(+)
|
|
|
|
diff --git a/tensorflow/core/kernels/maxpooling_op.cc b/tensorflow/core/kernels/maxpooling_op.cc
|
|
index 003d2e94b99cd..ceb6694ed665d 100644
|
|
--- a/tensorflow/core/kernels/maxpooling_op.cc
|
|
+++ b/tensorflow/core/kernels/maxpooling_op.cc
|
|
@@ -1088,6 +1088,8 @@ class MaxPoolingGradWithArgmaxOp : public OpKernel {
|
|
OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
|
|
{0}, 0, out_shape, &grad_out));
|
|
|
|
+ if (out_shape.num_elements() == 0) return; // nothing to be done
|
|
+
|
|
LaunchMaxPoolingGradWithArgmax<Device, T>::launch(
|
|
context, params, grad_in, argmax, grad_out, include_batch_in_index_);
|
|
}
|