import torch


def load_model(file_name, gpu=False):
    """Reads model from a file."""
    with open(file_name, 'rb') as f:
        model = torch.load(f, map_location=lambda storage, loc: storage)
        print('gpu:', gpu)
        if not gpu:
            model.set_device_id(None)
        else:
            model.cuda()
        return model