from typing import List
import torch


def average_checkpoints(
        checkpoint_paths: List[str],
):
    """
    Average a list of checkpoints' state_dicts.

    Parameters
    ----------
    checkpoint_paths
        A list of model checkpoint paths.

    Returns
    -------
    The averaged state_dict.
    """
    avg_state_dict = {}
    for per_path in checkpoint_paths:
        state_dict = torch.load(per_path, map_location=torch.device("cpu"))["state_dict"]
        for key in state_dict:
            if key in avg_state_dict:
                avg_state_dict[key] += state_dict[key]
            else:
                avg_state_dict[key] = state_dict[key]
        del state_dict

    num = torch.tensor(len(checkpoint_paths))
    for key in avg_state_dict:
        avg_state_dict[key] = avg_state_dict[key] / num.to(avg_state_dict[key])

    return avg_state_dict