def patch_load_state_dict(optimizer):
    r"""
    Re-pack parameter state after load state_dict
    """

    def repack(self):
        for group in self.param_groups:
            for _, p in enumerate(group["params"]):
                if p in self.params_attr and p.device.type == "cpu":
                    attr = self.params_attr[p]
                    if attr.op_ctx is not None:
                        # weight attr need "op" info to pack state while bias attr not
                        state = self.state[p]
                        plain_format_shape = attr.plain_format_shape
                        for state_key, state_value in state.items():
                            if (
                                isinstance(state_value, torch.Tensor)
                                and state_value.size() == plain_format_shape
                            ):
                                # We have an assumption here that any tensor's in parameter state, if they
                                # have same shapes with the parameter, they should share same layout with
                                # the parameter. Thus we need pack the state as we did to parameters.
                                pack_state(state, state_key, state_value, attr)

    def original_load_state_dict_without_state_cast(self, state_dict):
        r"""Loads the optimizer state.

        Args:
            state_dict (dict): optimizer state. Should be an object returned
                from a call to :meth:`state_dict`.
        Copied from torch/optim/optimizer.py.
        We need copy it here and change the behavior of state cast.
        For example, in out bf16 training. The mumentum buffer should always
        be float, but the original load_state_dict for optimizer will cast it to
        bfloat16 which will loss accuracy

        The original code:

            def cast(param, value, key=None):
                if isinstance(value, torch.Tensor):
                    # Floating-point types are a bit special here. They are the only ones
                    # that are assumed to always match the type of params.
                    # Make sure state['step'] is not casted https://github.com/pytorch/pytorch/issues/74424
                    if (key != "step"):
                        if param.is_floating_point():
                            value = value.to(param.dtype)
                        value = value.to(param.device)
                    return value
                elif isinstance(value, dict):
                    return {k: cast(param, v, key=k) for k, v in value.items()}
                elif isinstance(value, container_abcs.Iterable):
                    return type(value)(cast(param, v) for v in value)
                else:
                    return value
            state = defaultdict(dict)
            for k, v in state_dict['state'].items():
                if k in id_map:
                    param = id_map[k]
                    state[param] = cast(param, v)
                else:
                    state[k] = v

        We change it to:

            state = defaultdict(dict)
            for k, v in state_dict['state'].items():
                if k in id_map:
                    param = id_map[k]
                    state[param] = v
                else:
                    state[k] = v
        """
        # deepcopy, to be consistent with module API
        state_dict = deepcopy(state_dict)
        # Validate the state_dict
        groups = self.param_groups
        saved_groups = state_dict["param_groups"]

        if len(groups) != len(saved_groups):
            raise ValueError(
                "loaded state dict has a different number of " "parameter groups"
            )
        param_lens = (len(g["params"]) for g in groups)
        saved_lens = (len(g["params"]) for g in saved_groups)
        if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)):
            raise ValueError(
                "loaded state dict contains a parameter group "
                "that doesn't match the size of optimizer's group"
            )

        # Update the state
        id_map = dict(
            zip(
                chain.from_iterable((g["params"] for g in saved_groups)),
                chain.from_iterable((g["params"] for g in groups)),
            )
        )

        # Copy state assigned to params (and cast tensors to appropriate types).
        # State that is not assigned to params is copied as is (needed for
        # backward compatibility).
        state = defaultdict(dict)
        for k, v in state_dict["state"].items():
            if k in id_map:
                param = id_map[k]
                state[param] = v
            else:
                state[k] = v

        # Update parameter groups, setting their 'params' value
        def update_group(group, new_group):
            new_group["params"] = group["params"]
            return new_group

        param_groups = [update_group(g, ng) for g, ng in zip(groups, saved_groups)]
        self.__setstate__({"state": state, "param_groups": param_groups})

    def load_state_dict(self, state_dict):
        original_load_state_dict_without_state_cast(self, state_dict)
        repack(self)

    if not hasattr(optimizer, "_original_load_state_dict"):
        setattr(  # noqa: B010
            optimizer, "_original_load_state_dict", optimizer.load_state_dict
        )
        optimizer.load_state_dict = types.MethodType(load_state_dict, optimizer)

    Re-pack parameter state after load state_dict
    
intel_extension_for_pytorch/optim/_optimizer_utils.py