Don't do any work if output tensor is null (prevent div by 0)

PiperOrigin-RevId: 372208700 Change-Id: Iea6b6293e887ade8538facfdb50fb931e17f511e
tensorflow · May 5, 2021 · 376c352 · 376c352
1 parent 279bab6
commit 376c352
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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_);
   }