import torch


def load_model_from_cwd(**kwargs):
    """
    Loads structure and best model params from cwd, returns model only.
    :param kwargs: passed to torch.load

    :return:
    """

    with open("experiment_structure.pt", "rb") as pfile:
        structure = torch.load(pfile, **kwargs)

    with open("best_model.pkl", "rb") as pfile:
        state = torch.load(pfile, **kwargs)

    model = structure["training_modules"].model
    model.load_state_dict(state)

    return model