Skip to content
Permalink
Browse files Browse the repository at this point in the history
Add one missing valdiation to matrix_set_diag_op.cc
PiperOrigin-RevId: 387923408
Change-Id: If6a97b9098c13879400f56c22f91555cdf0ce5d7
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Jul 31, 2021
1 parent 59077ac commit ff88940
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions tensorflow/core/kernels/linalg/matrix_set_diag_op.cc
Expand Up @@ -70,6 +70,9 @@ class MatrixSetDiagOp : public OpKernel {
errors::InvalidArgument(
"diag_index must be a scalar or vector, received shape: ",
diag_index.shape().DebugString()));
OP_REQUIRES(
context, diag_index.NumElements() > 0,
errors::InvalidArgument("diag_index must have at least one element"));
lower_diag_index = diag_index.flat<int32>()(0);
upper_diag_index = lower_diag_index;
if (TensorShapeUtils::IsVector(diag_index.shape())) {
Expand Down

0 comments on commit ff88940

Please sign in to comment.