From 47f21101fa24bc70424f0330752447819669172c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 21 Oct 2020 14:36:04 -0700
Subject: [PATCH] Enable saved model restoration using just a .pb file.
PiperOrigin-RevId: 338342231
Change-Id: Idfb9a0692fb414d8bc153caf031fe9987fa83b30
---
tensorflow/cc/saved_model/loader.cc | 10 ++++++----
1 file changed, 6 insertions(+), 4 deletions(-)
diff --git a/tensorflow/cc/saved_model/loader.cc b/tensorflow/cc/saved_model/loader.cc
index 70d080a682f..dcd652d9fdf 100644
--- a/tensorflow/cc/saved_model/loader.cc
+++ b/tensorflow/cc/saved_model/loader.cc
@@ -404,10 +404,12 @@ Status RestoreSession(const RunOptions& run_options,
const uint64 read_start_microseconds = Env::Default()->NowMicros();
std::vector<AssetFileDef> asset_file_defs;
TF_RETURN_IF_ERROR(internal::GetAssetFileDefs(meta_graph, &asset_file_defs));
- TF_RETURN_IF_ERROR(RunRestore(run_options, export_dir,
- meta_graph.saver_def().restore_op_name(),
- meta_graph.saver_def().filename_tensor_name(),
- asset_file_defs, session->get()));
+ if (meta_graph.has_saver_def()) {
+ TF_RETURN_IF_ERROR(RunRestore(run_options, export_dir,
+ meta_graph.saver_def().restore_op_name(),
+ meta_graph.saver_def().filename_tensor_name(),
+ asset_file_defs, session->get()));
+ }
// Record walltime spent in restoring graph from disk, but postpone metric
// increments until graph init finishes.
const uint64 restore_graph_walltime =