chore: fix pipeline test

zama-ai · Sep 14, 2023 · 5cd1cc6 · 5cd1cc6
1 parent 74e5e24
commit 5cd1cc6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/sklearn/test_sklearn_models.py b/tests/sklearn/test_sklearn_models.py
@@ -807,7 +807,8 @@ def get_hyper_param_combinations(model_class):
             "base_score": [0.5, None],
         }
     elif model_class in get_sklearn_neighbors_models():
-        hyper_param_combinations = {"n_neighbors": [2, 4]}
+        # Use small `n_neighbors` values for KNN, because the data-set is too small for now
+        hyper_param_combinations = {"n_neighbors": [2]}
     else:
 
         assert is_model_class_in_a_list(