def _fix_projection_head_keys(state_dict):
    """Makes the state_dict compatible with the refactored projection heads.

    TODO: Remove once the models are refactored and the old checkpoints were
    replaced! Relevant issue: https://github.com/lightly-ai/lightly/issues/379

    Prevents unexpected key error when loading old checkpoints.

    """

    projection_head_identifier = "projection_head"
    prediction_head_identifier = "prediction_head"
    projection_head_insert = "layers"

    new_state_dict = {}
    for key, item in state_dict.items():
        if (
            projection_head_identifier in key or prediction_head_identifier in key
        ) and projection_head_insert not in key:
            # insert layers if it's not part of the key yet
            key_parts = key.split(".")
            key_parts.insert(1, projection_head_insert)
            new_key = ".".join(key_parts)
        else:
            new_key = key

        new_state_dict[new_key] = item

    return new_state_dict
Makes the state_dict compatible with the refactored projection heads.

    TODO: Remove once the models are refactored and the old checkpoints were
    replaced! Relevant issue: https://github.com/lightly-ai/lightly/issues/379

    Prevents unexpected key error when loading old checkpoints.

    
lightly/cli/_helpers.py