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 _save_model(self, *args, **kwargs):
      optimizer = self._get_optimizer()
      assert isinstance(optimizer, AveragedOptimizerWrapper)
  
      if self.update_weights:
          optimizer.assign_average_vars(self.model.trainable_weights)
          return super()._save_model(*args, **kwargs)
      else:
          # Note: `model.get_weights()` gives us the weights (non-ref)
          # whereas `model.variables` returns references to the variables.
          non_avg_weights = self.model.get_weights()
          optimizer.assign_average_vars(self.model.trainable_weights)
          # result is currently None, since `super._save_model` doesn't
          # return anything, but this may change in the future.
          result = super()._save_model(*args, **kwargs)
          self.model.set_weights(non_avg_weights)
          return result
  
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 4 for calling the AverageModelCheckpoint obj's underlying member method code with example usage: obj._save_model() and returns: super, result