keras/src/callbacks/swap_ema_weights.py/SwapEMAWeights/on_epoch_end
class SwapEMAWeights:    def on_epoch_end(self, epoch, logs=None):
        if self.swap_on_epoch and not self._ema_weights_in_model:
            self._swap_variables()
            self._ema_weights_in_model = True
            # We need to recover EMA weights from the previously swapped weights
            # in the last epoch. This is because, at the end of the fitting,
            # `finalize_variable_values` will be called to assign
            # `_model_variables_moving_average` to `trainable_variables`.
            if epoch == self.params["epochs"] - 1:
                self._finalize_ema_values()