import torch
from typing import OrderedDict


def load_torch_model(model, model_path):
    """Load state dict to model.

    Args:
        model: model to be loaded
        model_path: state dict file path

    Returns:
        loaded model
    """
    pretrained_model_dict = torch.load(model_path)
    new_state_dict = OrderedDict()
    for k, value in pretrained_model_dict.items():
        # name = k[7:]  # remove `module.`
        new_state_dict[k] = value
    model.load_state_dict(new_state_dict, strict=True)
    return model