From f1fee6e8ec4d98abb74f0347a911fea6cb69d0e8 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 cd9680c8212f6c..884c019dad5f90 100644
--- a/tensorflow/core/kernels/maxpooling_op.cc
+++ b/tensorflow/core/kernels/maxpooling_op.cc
@@ -1077,6 +1077,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_);
   }