from typing import Dict
from typing import Any


def convert_avg_params(checkpoint: Dict[str, Any]) -> Dict[str, Any]:
    """Converts average state dict to the format that can be loaded to a model.
    Args:
        checkpoint: model.
    Returns:
        Converted average state dict.
    """
    state_dict = checkpoint["state_dict"]
    avg_weights = {}
    for k, v in state_dict.items():
        if "avg_model" in k:
            avg_weights[k[10:]] = v
    return avg_weights