from typing import List
from pathlib import Path
import torch


def average_checkpoints(filenames: List[Path]) -> dict:
    """Average a list of checkpoints.

    Args:
      filenames:
        Filenames of the checkpoints to be averaged. We assume all
        checkpoints are saved by :func:`save_checkpoint`.
    Returns:
      Return a dict (i.e., state_dict) which is the average of all
      model state dicts contained in the checkpoints.
    """
    n = len(filenames)

    avg = torch.load(filenames[0], map_location="cpu")["model"]
    for i in range(1, n):
        state_dict = torch.load(filenames[i], map_location="cpu")["model"]
        for k in avg:
            avg[k] += state_dict[k]

    for k in avg:
        if avg[k].is_floating_point():
            avg[k] /= n
        else:
            avg[k] //= n

    return avg