class_docstr: The callback that saves average model weights.
  
  The callback that should be used with optimizers that extend
  `tfa.optimizers.AveragedOptimizerWrapper`, i.e.,
  `tfa.optimizers.MovingAverage` and
  `tfa.optimizers.StochasticAverage` optimizers.
  It saves and, optionally, assigns the averaged weights.
  
  Args:
      update_weights: If `True`, assign the moving average weights
          to the model, and save them. If False, keep the old
          non-averaged weights, but the saved model uses the
          average weights.
  
      See `tf.keras.callbacks.ModelCheckpoint` for the other args.
code: def set_model(self, model):
      super().set_model(model)
      optimizer = self._get_optimizer()
      if not isinstance(optimizer, AveragedOptimizerWrapper):
          raise TypeError(
              "AverageModelCheckpoint is only used when training"
              "with MovingAverage or StochasticAverage"
          )
  
file_path: addons/tensorflow_addons/callbacks/average_model_checkpoint.py
imports: import tensorflow
  from typeguard import typechecked
  from tensorflow_addons.optimizers.average_wrapper import AveragedOptimizerWrapper
label_desc: Use image node_id 3 for calling the AverageModelCheckpoint obj's underlying member method code with example usage: obj.set_model(model) without return types