36 lines
1.5 KiB
Diff
36 lines
1.5 KiB
Diff
From 9e82dce6e6bd1f36a57e08fa85af213e2b2f2622 Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
Date: Mon, 2 Aug 2021 14:21:41 -0700
|
|
Subject: [PATCH] Fix NPE in restoring code.
|
|
|
|
PiperOrigin-RevId: 388303253
|
|
Change-Id: Ia8c68568cb854bca538909a182b31a618d68ce55
|
|
---
|
|
tensorflow/core/kernels/save_restore_tensor.cc | 9 ++++++++-
|
|
1 file changed, 8 insertions(+), 1 deletion(-)
|
|
|
|
diff --git a/tensorflow/core/kernels/save_restore_tensor.cc b/tensorflow/core/kernels/save_restore_tensor.cc
|
|
index 953c1dfb6290b..dcbed428a5a5a 100644
|
|
--- a/tensorflow/core/kernels/save_restore_tensor.cc
|
|
+++ b/tensorflow/core/kernels/save_restore_tensor.cc
|
|
@@ -151,11 +151,18 @@ void RestoreTensor(OpKernelContext* context,
|
|
context, size == 1,
|
|
errors::InvalidArgument(
|
|
"Input 0 (file_pattern) must be a string scalar; got a tensor of ",
|
|
- size, "elements"));
|
|
+ size, " elements"));
|
|
}
|
|
const string& file_pattern = file_pattern_t.flat<tstring>()(0);
|
|
|
|
const Tensor& tensor_name_t = context->input(1);
|
|
+ {
|
|
+ const int64_t size = tensor_name_t.NumElements();
|
|
+ OP_REQUIRES(context, size > restore_index,
|
|
+ errors::InvalidArgument(
|
|
+ "Input 1 (file_pattern) must be a have at least ",
|
|
+ restore_index + 1, " elements"));
|
|
+ }
|
|
const string& tensor_name = tensor_name_t.flat<tstring>()(restore_index);
|
|
|
|
// If we cannot find a cached reader we will allocate our own.
|