class_docstr: Keras callback which updates pruning wrappers with the optimizer step.
  
  This callback must be used when training a model which needs to be pruned. Not
  doing so will throw an error.
  
  Example:
  
  ```python
  model.fit(x, y,
      callbacks=[UpdatePruningStep()])
  ```
class_docstr_tok: ["Keras","callback","which","updates","pruning","wrappers","with","the","optimizer","step",".","This","callback","must","be","used","when","training","a","model","which","needs","to","be","pruned",".","Not","doing","so","will","throw","an","error",".","Example",":","``","`","python","model.fit","(","x",",","y",",","callbacks=","[","UpdatePruningStep","(",")","]",")","``","`"]
code: def on_epoch_end(self, batch, logs=None):
      # At the end of every epoch, remask the weights. This ensures that when
      # the model is saved after completion, the weights represent mask*weights.
      weight_mask_ops = []
  
      for layer in self.prunable_layers:
          if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
              if tf.executing_eagerly():
                  layer.pruning_obj.weight_mask_op()
              else:
                  weight_mask_ops.append(
                      layer.pruning_obj.weight_mask_op()
                  )
  
      K.batch_get_value(weight_mask_ops)
  
file_path: model-optimization/tensorflow_model_optimization/python/core/sparsity/keras/pruning_callbacks.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import six
  import tensorflow
  from tensorflow_model_optimization.python.core.keras import compat
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper
label_desc: Use image node_id 3 for calling the UpdatePruningStep obj's underlying member method code with example usage: obj.on_epoch_end(batch, logs) without return types