Use nonexperimental LSO API in base_task.py.

This shouldn't break any official models, since I changed all LossScaleOptimizer isinstance checks to use the nonexperimental version (the experimental LSO subclasses the nonexperimental LSO, so changing isinstance checks in this way is always safe).

PiperOrigin-RevId: 366891847

official/core/base_task.py

@@ -81,7 +81,7 @@ class Task(tf.Module, metaclass=abc.ABCMeta):
           optimizer,
           use_float16=runtime_config.mixed_precision_dtype == "float16",
           loss_scale=runtime_config.loss_scale,
-          use_experimental_api=True)
+          use_experimental_api=False)
 
     return optimizer
 

official/vision/beta/projects/yolo/tasks/image_classification.py

@@ -86,16 +86,14 @@ class ImageClassificationTask(image_classification.ImageClassificationTask):
 
       # For mixed_precision policy, when LossScaleOptimizer is used, loss is
       # scaled for numerical stability.
-      if isinstance(
-          optimizer, tf.keras.mixed_precision.experimental.LossScaleOptimizer):
+      if isinstance(optimizer, tf.keras.mixed_precision.LossScaleOptimizer):
         scaled_loss = optimizer.get_scaled_loss(scaled_loss)
 
     tvars = model.trainable_variables
     grads = tape.gradient(scaled_loss, tvars)
     # Scales back gradient before apply_gradients when LossScaleOptimizer is
     # used.
-    if isinstance(
-        optimizer, tf.keras.mixed_precision.experimental.LossScaleOptimizer):
+    if isinstance(optimizer, tf.keras.mixed_precision.LossScaleOptimizer):
       grads = optimizer.get_unscaled_gradients(grads)
 
     # Apply gradient clipping.

official/vision/beta/projects/yt8m/tasks/yt8m_task.py

@@ -184,7 +184,7 @@ class YT8MTask(base_task.Task):
       # For mixed_precision policy, when LossScaleOptimizer is used, loss is
       # scaled for numerical stability.
       if isinstance(optimizer,
-                    tf.keras.mixed_precision.experimental.LossScaleOptimizer):
+                    tf.keras.mixed_precision.LossScaleOptimizer):
         scaled_loss = optimizer.get_scaled_loss(scaled_loss)
 
     tvars = model.trainable_variables
@@ -192,7 +192,7 @@ class YT8MTask(base_task.Task):
     # Scales back gradient before apply_gradients when LossScaleOptimizer is
     # used.
     if isinstance(optimizer,
-                  tf.keras.mixed_precision.experimental.LossScaleOptimizer):
+                  tf.keras.mixed_precision.LossScaleOptimizer):
       grads = optimizer.get_unscaled_gradients(grads)
 
     # Apply gradient clipping.