Temporary fix nightly package breakage caused by TF-MOT

PiperOrigin-RevId: 433540917

official/core/actions.py

@@ -21,7 +21,6 @@ from absl import logging
 import gin
 import orbit
 import tensorflow as tf
-import tensorflow_model_optimization as tfmot
 
 from official.core import base_trainer
 from official.core import config_definitions
@@ -52,6 +51,8 @@ class PruningAction:
       optimizer: `tf.keras.optimizers.Optimizer` optimizer instance used for
         training. This will be used to find the current training steps.
     """
+    # TODO(b/221490190): Avoid local import when the bug is fixed.
+    import tensorflow_model_optimization as tfmot  # pylint: disable=g-import-not-at-top
     self._optimizer = optimizer
     self.update_pruning_step = tfmot.sparsity.keras.UpdatePruningStep()
     self.update_pruning_step.set_model(model)