No public description

PiperOrigin-RevId: 569394081

Author: baek-jinoo

official/nlp/tasks/question_answering_test.py

@@ -101,14 +101,16 @@ def _run_task(self, config):
     logs = task.aggregate_logs(step_outputs=logs)
     metrics = task.reduce_aggregated_logs(logs)
     self.assertIn("final_f1", metrics)
-    model.save(os.path.join(self.get_temp_dir(), "saved_model"))
+    model.save(os.path.join(self.get_temp_dir(), "saved_model.keras"),
+               save_format="keras")
 
   @parameterized.parameters(
       itertools.product(
           (False, True),
           ("WordPiece", "SentencePiece"),
       ))
   def test_task(self, version_2_with_negative, tokenization):
+    del tokenization
     # Saves a checkpoint.
     pretrain_cfg = bert.PretrainerConfig(
         encoder=self._encoder_config,