class_docstr: _save_model from AverageModelCheckpoint is missing the batch arg.
code: def on_test_end(self, logs):
      if self.update_weights:
          self._get_optimizer().swap_weights()
  
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 4 for calling the AverageModelCheckpointPatched obj's underlying member method code with example usage: obj.on_test_end(logs) without return types