Optimizer change to get compatible for an incoming Keras optimizer migration.

PiperOrigin-RevId: 477217857

official/projects/mobilebert/distillation.py

@@ -238,6 +238,9 @@ class BertDistillationTask(policies.ProgressivePolicy, base_task.Task):
             })
     opt_factory = optimization.OptimizerFactory(params)
     optimizer = opt_factory.build_optimizer(opt_factory.build_learning_rate())
+    if isinstance(optimizer, tf.keras.optimizers.experimental.Optimizer):
+      optimizer = tf.keras.__internal__.optimizers.convert_to_legacy_optimizer(
+          optimizer)
 
     return optimizer
 

official/projects/mobilebert/distillation_test.py

@@ -153,7 +153,7 @@ class DistillationTest(tf.test.TestCase, parameterized.TestCase):
 
     eval_dataset = bert_distillation_task.get_eval_dataset(stage_id=0)
     eval_iterator = iter(eval_dataset)
-    optimizer = tf.keras.optimizers.SGD(lr=0.1)
+    optimizer = tf.keras.optimizers.legacy.SGD(learning_rate=0.1)
 
     # test train/val step for all stages, including the last pretraining stage
     for stage in range(student_block_num + 1):