Merge pull request #47221 from around-star:fix-avgpool-gpu

PiperOrigin-RevId: 359057811 Change-Id: I29cd1453f2360f9bd109bb6de0844a3baf06a66a
tensorflow · Feb 23, 2021 · fb21da8 · fb21da8
2 parents 493b780 + 83604e5
commit fb21da8
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/avgpooling_op.cc b/tensorflow/core/kernels/avgpooling_op.cc
@@ -152,6 +152,11 @@ class AvgPoolingOp<GPUDevice, T> : public UnaryOp<T> {
     OP_REQUIRES(context, ksize_n == 1 && stride_n == 1,
                 errors::Unimplemented(
                     "Pooling is not yet supported on the batch dimension."));
+
+    for (int i = 0; i < ksize_.size(); ++i) {
+      OP_REQUIRES(context, ksize_[i] != 0,
+                  errors::InvalidArgument("ksize cannot be zero"));
+    }
   }
 
   void Compute(OpKernelContext* context) override {