def patch_state_dict(optimizer):
    r"""
    To support resume training.
    Patch "state_dict" method to return unpacked/FP32 parameters/states
    """

    def get_optimizer_unpacked_state_dict(self):
        opt = self
        opt_temp = copy.deepcopy(opt)
        for (k1, _), (_, v2) in zip(opt.state.items(), opt_temp.state.items()):
            if k1 in opt.params_attr:
                attr = opt.params_attr[k1]
                for state_key, state_value in v2.items():
                    if (
                        isinstance(state_value, torch.Tensor)
                        and state_value.shape == k1.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 unpack the state as we did to parameters.
                        if attr.op_ctx is not None:
                            assert attr.num_modules == 1
                            state_value = attr.op_ctx.to_public(state_value)
                        v2[state_key] = state_value
        return opt_temp.state_dict()

    if not hasattr(optimizer, "_original_state_dict"):
        setattr(optimizer, "_original_state_dict", optimizer.state_dict)  # noqa: B010
        optimizer.state_dict = types.MethodType(
            get_optimizer_unpacked_state_dict, optimizer
        )

    To support resume training.
    Patch "state_dict" method to return unpacked/FP32 parameters/states
    
intel_extension_for_pytorch/optim/_optimizer_utils.py