add params for xgboost

euxhenh · Dec 9, 2023 · 317dc63 · 317dc63
1 parent 2cfab33
commit 317dc63
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/tests/test_predictors.py b/tests/test_predictors.py
@@ -148,13 +148,22 @@ def test_classifiers_x(X, classifier, key):
     kmeans = cfg_kmeans.create()
     kmeans(adata)
 
+    if classifier == "XGBClassifier":
+        kwargs = {
+            'n_estimators': 1,
+            'max_depth': 1,
+        }
+    else:
+        kwargs = {}
+
     cfg = OmegaConf.create(
         {
             "_target_": f"src.grinch.{classifier}.Config",
             "x_key": f"obsm.{OBSM.X_PCA}",
             "y_key": f"obs.{OBS.KMEANS}",
             "seed": 42,
             "labels_key": f"obs.{key}",
+            **kwargs,
         }
     )
     # Need to start using convert all for lists and dicts