class_docstr: _save_model from AverageModelCheckpoint is missing the batch arg.
code: def _save_model(self, epoch, batch, logs):
      optimizer = self._get_optimizer()
      assert isinstance(optimizer, AveragedOptimizerWrapper)
  
      if self.update_weights:
          optimizer.assign_average_vars(self.model.variables)
  
          # Use super call two levels up to bypass the problematic function
          return super(AverageModelCheckpoint, self)._save_model(
              epoch, batch, logs
          )
      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.variables)
          # result is currently None, since `super._save_model` doesn't
          # return anything, but this may change in the future.
  
          # Use call two levels up to bypass the problematic function
          result = super(AverageModelCheckpoint, self)._save_model(
              epoch, batch, logs
          )
          self.model.set_weights(non_avg_weights)
          return result
  
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 2 for calling the AverageModelCheckpointPatched obj's underlying member method code with example usage: obj._save_model(epoch, batch, logs) and returns: super, result