def METHOD_NAME(self, state_dict=None):
    """
    Store a copy of the EMA params in fp32.
    If state dict is passed, the EMA params is copied from
    the provided state dict. Otherwise, it is copied from the
    current EMA model parameters.
    """
    if not self.config.ema_fp32:
        raise RuntimeError(
            "build_fp32_params should not be called if ema_fp32=False. "
            "Use ema_fp32=True if this is really intended."
        )
    if state_dict is None:
        state_dict = self.model.state_dict()
    def _to_float(t):
        return t.float() if torch.is_floating_point(t) else t
    for param_key in state_dict:
        if param_key in self.fp32_params:
            if param_key == "__sq_mom":
                self.fp32_params[param_key] = state_dict[param_key]
            else:
                self.fp32_params[param_key].copy_(state_dict[param_key])
        else:
            self.fp32_params[param_key] = _to_float(state_dict[param_key])
            if "__sq_mom" in self.fp32_params:
                self.fp32_params["__sq_mom"][param_key] = torch.zeros_like(
                    self.fp32_params[param_key]
                )