def _rekey_optimizer_state_if_needed(optimizer_state_dict: dict[str, Any], module: Module) -> dict[str, Any]:
    """Handles the case where the optimizer state is saved from a normal optimizer and converts the keys to parameter
    names."""
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
    from torch.distributed.fsdp import OptimStateKeyType

    if isinstance(list(optimizer_state_dict["state"].keys())[0], int):
        optimizer_state_dict = FSDP.rekey_optim_state_dict(optimizer_state_dict, OptimStateKeyType.PARAM_NAME, module)
    return optimizer_state_dict
Handles the case where the optimizer state is saved from a normal optimizer and converts the keys to parameter
    names.
src/lightning/fabric/strategies/model_parallel.py