def METHOD_NAME(self, new_model, updates=None):
    """ One update of the EMA model based on new model weights """
    decay = self.decay
    ema_state_dict = {}
    ema_params = self.fp32_params if self.config.ema_fp32 else self.model.state_dict()
    for key, param in new_model.state_dict().items():
        try:
            ema_param = ema_params[key]
        except KeyError:
            ema_param = param.float().clone() if param.ndim == 1 else copy.deepcopy(param)
        if param.shape != ema_param.shape:
            raise ValueError(
                "incompatible tensor shapes between model param and ema param"
                + "{} vs. {}".format(param.shape, ema_param.shape)
            )
        if "version" in key:
            # Do not decay a model.version pytorch param
            continue
        ema_param.mul_(decay)
        ema_param.add_(param.to(dtype=ema_param.dtype), alpha=1-decay)
        ema_state_dict[key] = ema_param
    self.restore(ema_state_dict, build_fp32_params=False)