import torch


def save_model(MainModel, network_filepath, weight_filepath, dump_filepath):
    model = MainModel.KitModel(weight_filepath)
    model.eval()
    torch.save(model, dump_filepath)
    print('PyTorch model file is saved as [{}], generated by [{}.py] and [{}]. Notice that you may need [{}.py] to load the model back.'.format(
        dump_filepath, network_filepath, weight_filepath, network_filepath))
