Prevent overflow due to integer conversion to unsigned.

PiperOrigin-RevId: 387738045 Change-Id: Id7e95bc07e02df1c66b72bd09f389608c87bdebe
tensorflow · Jul 30, 2021 · b5cdbf1 · b5cdbf1
1 parent 532f5c5
commit b5cdbf1
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/image/non_max_suppression_op.cc b/tensorflow/core/kernels/image/non_max_suppression_op.cc
@@ -930,6 +930,8 @@ class CombinedNonMaxSuppressionOp : public OpKernel {
         errors::InvalidArgument("max_size_per_class must be 0-D, got shape ",
                                 max_output_size.shape().DebugString()));
     const int max_size_per_class = max_output_size.scalar<int>()();
+    OP_REQUIRES(context, max_size_per_class > 0,
+                errors::InvalidArgument("max_size_per_class must be positive"));
     // max_total_size: scalar
     const Tensor& max_total_size = context->input(3);
     OP_REQUIRES(