Function remove_model_prefix

lit_gpt/utils.py:540–549

Handles the case where the LightningModule wrapper (FullPintsTransformer) prepends 'model.' to all keys in the state dict. Removes the prefix in-place.

(state_dict: Dict[str, torch.Tensor])

def remove_model_prefix(state_dict: Dict[str, torch.Tensor]) -> None:
    """
    Handles the case where the LightningModule wrapper (FullPintsTransformer) prepends 'model.'
    to all keys in the state dict. Removes the prefix in-place.
    """
    keys = list(state_dict.keys())
    for key in keys:
        new_key = key.removeprefix('model.')
        if new_key != key:
            state_dict[new_key] = state_dict.pop(key)


def add_model_prefix(state_dict: Dict[str, torch.Tensor]) -> None: