Code to soup two models in pytorch

soup.py

def soup_two_models(model, second_model):
    souped_model = copy.deepcopy(model)
    for param in souped_model.named_parameters():
        name = param[0]
        param[1].data = (model.state_dict()[name] + second_model.state_dict()[name]) / 2

    return souped_model