Merge pull request #51280 from tensorflow/mm-cherrypick-d7de67733925d…

…e196ec8863a33445b73f9562d1d-on-r2.3 Prevent a CHECK-fail due to empty tensor input in `map_stage_op.cc`
tensorflow · Aug 5, 2021 · a695e18 · a695e18
2 parents bf5ddbc + cc96b00
commit a695e18
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/map_stage_op.cc b/tensorflow/core/kernels/map_stage_op.cc
@@ -530,6 +530,8 @@ class MapStageOp : public OpKernel {
     OP_REQUIRES_OK(ctx, ctx->input("key", &key_tensor));
     OP_REQUIRES_OK(ctx, ctx->input("indices", &indices_tensor));
     OP_REQUIRES_OK(ctx, ctx->input_list("values", &values_tensor));
+    OP_REQUIRES(ctx, key_tensor->NumElements() > 0,
+                errors::InvalidArgument("key must not be empty"));
 
     // Create copy for insertion into Staging Area
     Tensor key(*key_tensor);