Skip to content

GitLab

Commit 6d82d204 authored by Hongkun Yu's avatar Hongkun Yu Committed by A. Unique TensorFlower
Browse files

Temporary fix nightly package breakage caused by TF-MOT 

PiperOrigin-RevId: 433540917
parent ce730af7
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
official/core/actions.py
View file @ 6d82d204
...@@ -21,7 +21,6 @@ from absl import logging ...@@ -21,7 +21,6 @@ from absl import logging
import gin import gin
import orbit import orbit
import tensorflow as tf import tensorflow as tf
import tensorflow_model_optimization as tfmot
from official.core import base_trainer from official.core import base_trainer
from official.core import config_definitions from official.core import config_definitions
...@@ -52,6 +51,8 @@ class PruningAction: ...@@ -52,6 +51,8 @@ class PruningAction:
optimizer: `tf.keras.optimizers.Optimizer` optimizer instance used for optimizer: `tf.keras.optimizers.Optimizer` optimizer instance used for
training. This will be used to find the current training steps. 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._optimizer = optimizer
self.update_pruning_step = tfmot.sparsity.keras.UpdatePruningStep() self.update_pruning_step = tfmot.sparsity.keras.UpdatePruningStep()
self.update_pruning_step.set_model(model) self.update_pruning_step.set_model(model)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment