From 21745f6a97cc5a88ec1749b166ddf002789bc645 Mon Sep 17 00:00:00 2001
From: Scott Zhu <scottzhu@google.com>
Date: Mon, 1 Mar 2021 13:23:21 -0800
Subject: [PATCH] Fix save model related unit test.

1. The save_test requires cached_session() in graph mode.
2. Callback test is failing in OSS the model output shape doesn't match with the input data. Not sure why its not failing in python/keras.

PiperOrigin-RevId: 360262858
Change-Id: Ie5943a83943d1d759f061c41bf8877c8fd6a0506
---
 tensorflow/python/keras/callbacks_test.py   |  2 +-
 tensorflow/python/keras/saving/save_test.py | 39 +++++++++++----------
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/tensorflow/python/keras/callbacks_test.py b/tensorflow/python/keras/callbacks_test.py
index bc5910a434d..a887e743c61 100644
--- a/tensorflow/python/keras/callbacks_test.py
+++ b/tensorflow/python/keras/callbacks_test.py
@@ -751,6 +751,7 @@ class KerasCallbacksTest(keras_parameterized.TestCase):
         test_samples=TEST_SAMPLES,
         input_shape=(INPUT_DIM,),
         num_classes=NUM_CLASSES)
+    y_train = np_utils.to_categorical(y_train, num_classes=NUM_CLASSES)
 
     model.fit(
         x_train,
@@ -760,7 +761,6 @@ class KerasCallbacksTest(keras_parameterized.TestCase):
         verbose=0)
     # Check that the filepath is a SavedModel directory.
     self.assertIn('saved_model.pb', os.listdir(filepath))
-    os.remove(filepath)
 
   def _get_dummy_resource_for_model_checkpoint_testing(self):
 
diff --git a/tensorflow/python/keras/saving/save_test.py b/tensorflow/python/keras/saving/save_test.py
index 69489e99061..7d8750c3e4b 100644
--- a/tensorflow/python/keras/saving/save_test.py
+++ b/tensorflow/python/keras/saving/save_test.py
@@ -1012,27 +1012,28 @@ class TestWholeModelSaving(keras_parameterized.TestCase):
     e.g. "head_0_accuracy" should not become "head_0_head_0_accuracy" after
     saving and loading a model.
     """
-    input_ = keras.Input((4,))
-    model = keras.Model(
-        input_,
-        [keras.layers.Softmax(name='head_0')(keras.layers.Dense(3)(input_)),
-         keras.layers.Softmax(name='head_1')(keras.layers.Dense(5)(input_))])
-    metric = keras.metrics.BinaryAccuracy()
-    model.compile(optimizer='rmsprop',
-                  loss='mse',
-                  metrics={'head_0': [metric, 'accuracy']})
+    with self.cached_session():
+      input_ = keras.Input((4,))
+      model = keras.Model(
+          input_,
+          [keras.layers.Softmax(name='head_0')(keras.layers.Dense(3)(input_)),
+           keras.layers.Softmax(name='head_1')(keras.layers.Dense(5)(input_))])
+      metric = keras.metrics.BinaryAccuracy()
+      model.compile(optimizer='rmsprop',
+                    loss='mse',
+                    metrics={'head_0': [metric, 'accuracy']})
 
-    # Run one iteration.
-    x = np.random.rand(2, 4)
-    y = {'head_0': np.random.randint(2, size=(2, 3)),
-         'head_1': np.random.randint(2, size=(2, 5))}
-    model.fit(x, y, verbose=0)
+      # Run one iteration.
+      x = np.random.rand(2, 4)
+      y = {'head_0': np.random.randint(2, size=(2, 3)),
+           'head_1': np.random.randint(2, size=(2, 5))}
+      model.fit(x, y, verbose=0)
 
-    # Save and reload.
-    save_format = testing_utils.get_save_format()
-    saved_model_dir = self._save_model_dir()
-    keras.models.save_model(model, saved_model_dir, save_format=save_format)
-    loaded = keras.models.load_model(saved_model_dir)
+      # Save and reload.
+      save_format = testing_utils.get_save_format()
+      saved_model_dir = self._save_model_dir()
+      keras.models.save_model(model, saved_model_dir, save_format=save_format)
+      loaded = keras.models.load_model(saved_model_dir)
 
     # Make sure the metrics names from the model before saving match the loaded
     # model.