Merge pull request #51186 from tensorflow/mm-cherrypick-9728c60e13691…

…2a12d99ca56e106b7cce7af5986-on-r2.4 Ensure validation sticks in `save_restore_v2_ops.cc`
tensorflow · Aug 5, 2021 · 4ca0469 · 4ca0469
2 parents 4de6435 + 3b13a5d
commit 4ca0469
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/save_restore_v2_ops.cc b/tensorflow/core/kernels/save_restore_v2_ops.cc
@@ -98,6 +98,7 @@ class SaveV2 : public OpKernel {
     const Tensor& shape_and_slices = context->input(2);
     ValidateInputs(true /* is save op */, context, prefix, tensor_names,
                    shape_and_slices);
+    if (!context->status().ok()) return;
 
     const int kFixedInputs = 3;  // Prefix, tensor names, shape_and_slices.
     const int num_tensors = static_cast<int>(tensor_names.NumElements());
@@ -177,6 +178,7 @@ class RestoreV2 : public OpKernel {
                                         " expected dtypes."));
     ValidateInputs(false /* not save op */, context, prefix, tensor_names,
                    shape_and_slices);
+    if (!context->status().ok()) return;
 
     const string& prefix_string = prefix.scalar<tstring>()();