class_docstr: _save_model from AverageModelCheckpoint is missing the batch arg.
code: def __init__(
      self,
      update_weights: bool,
      filepath: str,
      monitor: str = "val_loss",
      verbose: int = 0,
      save_best_only: bool = False,
      save_weights_only: bool = False,
      mode: str = "auto",
      save_freq: str = "epoch",
      best_metrics_from_checkpoint: float = float("-inf"),
      **kwargs,
  ):
      super().__init__(
          update_weights,
          filepath,
          monitor,
          verbose,
          save_best_only,
          save_weights_only,
          mode,
          save_freq,
          **kwargs,
      )
      if (
          best_metrics_from_checkpoint is not None
          and best_metrics_from_checkpoint > float("-inf")
      ):
          self.best = best_metrics_from_checkpoint
  
file_path: deepvariant/deepvariant/average_model_checkpoint_patched.py
imports: from tensorflow_addons.callbacks.average_model_checkpoint import AverageModelCheckpoint
  from tensorflow_addons.optimizers.average_wrapper import AveragedOptimizerWrapper
label_desc: Use image node_id 1 to create a new AverageModelCheckpointPatched object from inherited base classes: AverageModelCheckpoint with example: obj = AverageModelCheckpointPatched(update_weights, filepath, monitor, verbose, save_best_only, save_weights_only, mode, save_freq, best_metrics_from_checkpoint)