use gpu for xgboost

MaastrichtU-IDS · Dec 8, 2023 · 01692c6 · 01692c6
1 parent 2c014b8
commit 01692c6
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/train.py b/src/train.py
@@ -221,6 +221,7 @@ def train(
         objective='binary:logistic',  # For binary classification
         n_jobs=-1,
         random_state=42,
+        tree_method='gpu_hist', # Use GPU optimized histogram algorithm
     )
 
     # clfs = [('Naive Bayes',nb_model),('Logistic Regression',lr_model),('Random Forest',rf_model)]