import torch
from typing import OrderedDict


def load_statedict(model_path):
    """Loads model state dict.

    Args:
        model_path: model path

    Returns: state dict
    """
    print(f"Loading model from {model_path}.")
    state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
    print("Loaded model.")
    # Remove distributed naming if model trained in distributed mode
    model_state_dict = OrderedDict()
    for k, v in state_dict["model"].items():
        if k.startswith("module."):
            name = k[len("module.") :]
            model_state_dict[name] = v
        else:
            model_state_dict[k] = v
    return state_dict, model_state_dict