    def _step_internal(self, new_model):
        """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()
        )

        new_p = []
        ema_p = []

        for key, param in new_model.named_parameters():
            if isinstance(param, dict):
                continue

            if not self.add_missing_params and key not in ema_params:
                continue

            try:
                ema_param = ema_params[key]
            except KeyError:
                ema_param = (
                    param.float().clone() if param.ndim == 1 else copy.deepcopy(param)
                )
                ema_params[key] = ema_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

            lr = 1 - decay

            if key in self.skip_keys or not param.requires_grad:
                ema_params[key].copy_(param.to(dtype=ema_param.dtype).data)
                ema_param = ema_params[key]
            else:
                if self.log_norms:
                    new_p.append(param)
                    ema_p.append(ema_param)

                ema_param.mul_(1 - lr)
                ema_param.add_(param.data.to(dtype=ema_param.dtype), alpha=lr)

            ema_state_dict[key] = ema_param

        for key, param in new_model.named_buffers():
            ema_state_dict[key] = param

        if self.log_norms:
            if "model_norm" in self.logs:
                self.prev_model_norm = self.logs["model_norm"]

            chunk_size = 2048 * 32
            has_inf = torch.zeros(
                (1, 1), dtype=torch.int, device=next(new_model.parameters()).device
            )

            new_norm = multi_tensor_l2norm(chunk_size, has_inf, [new_p], False)
            old_norm = multi_tensor_l2norm(chunk_size, has_inf, [ema_p], False)

            self.logs["model_norm"] = new_norm[0]
            self.logs["ema_norm"] = old_norm[0]

        self.restore(ema_state_dict, build_fp32_params=False)
One update of the EMA model based on new model weights
fairseq/modules/ema_module.py