Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix overflow CHECK issue with tf.raw_ops.UnsortedSegmentJoin.
PiperOrigin-RevId: 370766155
Change-Id: I33e7c6626224e1060a8a4ab51ad5d861c6d4c63e
  • Loading branch information
Amit Patankar authored and tensorflower-gardener committed Apr 27, 2021
1 parent a2a607d commit 704866e
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions tensorflow/core/kernels/unsorted_segment_join_op.cc
Expand Up @@ -90,6 +90,8 @@ class UnsortedSegmentJoinOp : public OpKernel {
const int32 segment_dims = segment_id_shape.dims();

const Tensor& num_segments_tensor = context->input(2);
OP_REQUIRES(context, num_segments_tensor.NumElements() != 0,
errors::InvalidArgument("Number of segments cannot be empty."));
auto num_segments = num_segments_tensor.scalar<NUM_SEGMENTS_TYPE>()();

OP_REQUIRES(context, segment_dims != 0,
Expand Down

0 comments on commit 704866e

Please sign in to comment.