def set_models(self, pred_model, proba_model):
    """
        Sets the models to use
        :param pred_model: the prediction model
        :param proba_model: the model that outputs the activation maps
    """
    self.pred_model = pred_model
    self.proba_model = proba_model