Validate num_segments > 0 in unsorted_segment_join

Fixes #55305 PiperOrigin-RevId: 442047005
tensorflow · Apr 15, 2022 · 84563f2 · 84563f2
1 parent c718ee7
commit 84563f2
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/unsorted_segment_join_op.cc b/tensorflow/core/kernels/unsorted_segment_join_op.cc
@@ -94,6 +94,8 @@ class UnsortedSegmentJoinOp : public OpKernel {
                 errors::InvalidArgument("Number of segments cannot be empty."));
     auto num_segments = num_segments_tensor.scalar<NUM_SEGMENTS_TYPE>()();
 
+    OP_REQUIRES(context, num_segments > 0,
+                errors::InvalidArgument("Number of segments must be positive"));
     OP_REQUIRES(context, segment_dims != 0,
                 errors::InvalidArgument("Segment_id cannot have rank 0"));