From 45bc34f41a1092137e898c16193dc0d8cf37ce20 Mon Sep 17 00:00:00 2001 From: kcelia Date: Mon, 11 Sep 2023 15:06:25 +0200 Subject: [PATCH] chore: fix test_mono_param_waraning --- tests/sklearn/test_sklearn_models.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/sklearn/test_sklearn_models.py b/tests/sklearn/test_sklearn_models.py index aa45d837f4..6b2677ddb0 100644 --- a/tests/sklearn/test_sklearn_models.py +++ b/tests/sklearn/test_sklearn_models.py @@ -1461,11 +1461,10 @@ def test_predict_correctness( print("Compile the model") with warnings.catch_warnings(): - from concrete import fhe if get_model_name(model) == "KNeighborsClassifier": default_configuration.parameter_selection_strategy = ( - fhe.ParameterSelectionStrategy.MONO + ParameterSelectionStrategy.MONO ) fhe_circuit = model.compile( x, @@ -1690,6 +1689,10 @@ def test_mono_parameter_warnings( if is_model_class_in_a_list(model_class, get_sklearn_linear_models()): return + # KNN works only for ParameterSelectionStrategy.MULTI + if is_model_class_in_a_list(model_class, get_sklearn_neighbors_models()): + pytest.skip("Skipping predict_proba for KNN, doesn't work for now") + n_bits = min(N_BITS_REGULAR_BUILDS) model, x = preamble(model_class, parameters, n_bits, load_data, is_weekly_option)