Allow legacy optimizer in MG's optimizer factory.

PiperOrigin-RevId: 467067824
tensorflow · Aug 11, 2022 · 03c41a1 · 03c41a1
1 parent 433a87e
commit 03c41a1
Showing 1 changed file with 12 additions and 14 deletions.
diff --git a/official/modeling/optimization/optimizer_factory.py b/official/modeling/optimization/optimizer_factory.py
@@ -214,17 +214,15 @@ def build_optimizer(
           optimizer, **self._ema_config.as_dict())
     if postprocessor:
       optimizer = postprocessor(optimizer)
-    if not isinstance(optimizer, tf.keras.optimizers.Optimizer):
-      # tf.keras.optimizers.experimental only exist in tf-nightly.
-      # The following check makes sure the function wont' break in older TF
-      # version because of missing the experimental package.
-      if hasattr(tf.keras.optimizers, 'experimental'):
-        if not isinstance(optimizer,
-                          tf.keras.optimizers.experimental.Optimizer):
-          raise TypeError('OptimizerFactory.build_optimizer returning a '
-                          'non-optimizer object: {}'.format(optimizer))
-      else:
-        raise TypeError('OptimizerFactory.build_optimizer returning a '
-                        'non-optimizer object: {}'.format(optimizer))
-
-    return optimizer
+    if isinstance(optimizer, tf.keras.optimizers.Optimizer):
+      return optimizer
+    # The following check makes sure the function won't break in older TF
+    # version because of missing the experimental/legacy package.
+    if hasattr(tf.keras.optimizers, 'experimental'):
+      if isinstance(optimizer, tf.keras.optimizers.experimental.Optimizer):
+        return optimizer
+    if hasattr(tf.keras.optimizers, 'legacy'):
+      if isinstance(optimizer, tf.keras.optimizers.legacy.Optimizer):
+        return optimizer
+    raise TypeError('OptimizerFactory.build_optimizer returning a '
+                    'non-optimizer object: {}'.format(optimizer))