Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix FPE issue with tf.raw_ops.DenseCountSparseOutput.
PiperOrigin-RevId: 370946862
Change-Id: I3752584ad04aaecb327ff6793a9640ac56acfe7a
  • Loading branch information
Amit Patankar authored and tensorflower-gardener committed Apr 28, 2021
1 parent efff014 commit da5ff2d
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions tensorflow/core/kernels/count_ops.cc
Expand Up @@ -122,6 +122,9 @@ class DenseCount : public OpKernel {

int num_batch_elements = 1;
for (int i = 0; i < num_batch_dimensions; ++i) {
OP_REQUIRES(context, data.shape().dim_size(i) != 0,
errors::InvalidArgument(
"Invalid input: Shapes dimension cannot be 0."));
num_batch_elements *= data.shape().dim_size(i);
}
int num_value_elements = data.shape().num_elements() / num_batch_elements;
Expand Down

0 comments on commit da5ff2d

Please sign in to comment.