import torch
from typing import OrderedDict


def load_state_dict_from_multigpu2cpu(state_dir):
    """ To load weights trained using multiple gpu to cpu
    """
    multigpu_state_dict = torch.load(state_dir, map_location='cpu')
    new_state_dict = OrderedDict()
    for k, v in multigpu_state_dict.items():
        name = k[7:] # remove `module.`
        new_state_dict[name] = v
    return new_state_dict