import logging


def unwrap_state_dict(checkpoint, verbose=True):
    """Remove `module.` prefix for `state_dict`."""
    with_wrapper = False
    new_checkpoint = {}
    for key, val in checkpoint.items():
        if with_wrapper:  # all keys consistent
            assert key.startswith('module.')
        if key.startswith('module.'):
            _, new_key = key.split('.', maxsplit=1)
            if not with_wrapper and verbose:
                logging.info('Unwrap state_dict')
                with_wrapper = True
        else:
            new_key = key
        new_checkpoint[new_key] = val
    return new_checkpoint