keras/src/callbacks/swap_ema_weights.py/SwapEMAWeights/on_predict_begin
class SwapEMAWeights:    def on_predict_begin(self, logs=None):
        if not self._ema_weights_in_model:
            self._swap_variables()
            self._ema_weights_in_model = True