Skip to content
Permalink
Browse files Browse the repository at this point in the history
Ensure validation sticks in save_restore_v2_ops.cc
PiperOrigin-RevId: 387924206
Change-Id: I6156842eb3230076b5812c0815f3e66bd5241454
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Jul 31, 2021
1 parent f2a673b commit 9728c60
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions tensorflow/core/kernels/save_restore_v2_ops.cc
Expand Up @@ -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());
Expand Down Expand Up @@ -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>()();

Expand Down

0 comments on commit 9728c60

Please sign in to comment.