import torch
import pickle


def load_torch_model(file_name, history=False):
    """
    Loads a DLWPTorchNN or other model using Torch saved to disk with the `save_torch_model()` method.

    :param file_name: str: base name of save files
    :param history: bool: if True, loads the history file along with the model
    :return: model [, dict]: loaded object [, dictionary of training history]\
    """
    with open('%s.pkl' % file_name, 'rb') as f:
        model = pickle.load(f)
    model.model = torch.load('%s.torch' % file_name)
    model.model.eval()
    if history:
        with open('%s.history' % file_name, 'rb') as f:
            h = pickle.load(f)
        return model, h
    else:
        return model