def pack_optimizer_states(optimizer, param, attr):
    """
    1. convert user's optimizer weights and related states to packed format
    """
    if optimizer is None:
        return
    if param in optimizer.state:
        state = optimizer.state[param]
        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)

    1. convert user's optimizer weights and related states to packed format
    
intel_extension_for_pytorch/optim/_optimizer_utils.py