diff options
Diffstat (limited to 'tensorflow/core/kernels/segment_reduction_ops.cc')
-rw-r--r-- | tensorflow/core/kernels/segment_reduction_ops.cc | 3 |
1 files changed, 3 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/segment_reduction_ops.cc b/tensorflow/core/kernels/segment_reduction_ops.cc index 4302a68a18..2334e50f1d 100644 --- a/tensorflow/core/kernels/segment_reduction_ops.cc +++ b/tensorflow/core/kernels/segment_reduction_ops.cc @@ -376,6 +376,9 @@ struct UnsortedSegmentSumFunctor<CPUDevice, T, Index> auto data_flat = typename TTypes<T, 2>::ConstTensor(data, N, data_size / N); for (int64 i = 0; i < N; ++i) { Index j = internal::SubtleMustCopy(segment_ids(i)); + if (j < 0) { + continue; + } OP_REQUIRES(ctx, FastBoundsCheck(j, output_rows), errors::InvalidArgument( "segment_ids", SliceDebugString(segment_ids_shape, i), |