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 state_dict is None:
        state_dict = self.model.state_dict()
    # for non-float params (like registered symbols), they are copied into this dict and covered in each update
    for param_key in state_dict:
        if param_key in self.params:
            self.params[param_key].copy_(state_dict[param_key])
        else:
            self.params[param_key] = state_dict[param_key]