Use a sharded saver in SavedModel.

Change: 134471468

1 files changed, 14 insertions, 0 deletions

diff --git a/tensorflow/cc/saved_model/loader_test.cc b/tensorflow/cc/saved_model/loader_test.cc
index cef28d96f9..b2d55a9ade 100644
--- a/tensorflow/cc/saved_model/loader_test.cc
+++ b/tensorflow/cc/saved_model/loader_test.cc
@@ -28,6 +28,8 @@ namespace {
 
 constexpr char kTestDataPb[] = "cc/saved_model/testdata/half_plus_two";
 constexpr char kTestDataPbTxt[] = "cc/saved_model/testdata/half_plus_two_pbtxt";
+constexpr char kTestDataSharded[] =
+    "cc/saved_model/testdata/half_plus_two_sharded";
 
 class LoaderTest : public ::testing::Test {
  protected:
@@ -110,6 +112,18 @@ TEST_F(LoaderTest, PbtxtFormat) {
   CheckSavedModelBundle(bundle);
 }
 
+TEST_F(LoaderTest, ShardedVariables) {
+  SavedModelBundle bundle;
+  SessionOptions session_options;
+  RunOptions run_options;
+
+  const string export_dir =
+      io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
+  TF_ASSERT_OK(LoadSavedModel(export_dir, {kSavedModelTagServe},
+                              session_options, run_options, &bundle));
+  CheckSavedModelBundle(bundle);
+}
+
 TEST_F(LoaderTest, InvalidExportPath) {
   SavedModelBundle bundle;
   RunOptions run_options;