import torch


def load_checkpoint(checkpoint_path, gpu):
    """
    Load a saved model from a path.
    :returns: model
    """
    if gpu:
        checkpoint = torch.load(checkpoint_path)
    else:
        checkpoint = torch.load(
            checkpoint_path, map_location=lambda storage, loc: storage)
    model = checkpoint['model']
    model.class_to_idx = checkpoint['class_to_idx']
    model.load_state_dict(checkpoint['state_dict'])
    model.idx_to_class = dict([[v, k] for k, v in model.class_to_idx.items()])

    return model