def METHOD_NAME(
    self, batch: int, logs: Optional[Dict[str, Any]] = None
) -> None:
    """Called at the end of a training batch in `fit` methods."""
    self.on_batch_end(batch, logs if logs else {})