message: Use `SetDimWithStatus` instead of `set_dim`.
  This would prevent a crash due to overflow.
diff: @@ -20,6 +20,7 @@ limitations under the License.
  #include <cstdint>
  +#include "tensorflow/core/framework/op_requires.h"
  #include "tensorflow/core/platform/types.h"
  #define EIGEN_USE_THREADS
  #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  @@ -1100,7 +1101,7 @@ class SparseSegmentGradOpBase : public OpKernel {
  const auto segment_vec = segment_ids.vec<SegmentId>();
  TensorShape output_shape = input.shape();
  - output_shape.set_dim(0, M);
  + OP_REQUIRES_OK(context, output_shape.SetDimWithStatus(0, M));
  Tensor* output = nullptr;
  OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
  if (M == 0 || N == 0) return;
  