From b1870941f10bf30791f05460222013e26c0d5a39 Mon Sep 17 00:00:00 2001 From: dsethz Date: Sat, 9 Nov 2024 14:07:42 +0100 Subject: [PATCH] use all cpus --- src/nuclai/models/train_lr_rf.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/nuclai/models/train_lr_rf.py b/src/nuclai/models/train_lr_rf.py index c0952d0..3665e7d 100755 --- a/src/nuclai/models/train_lr_rf.py +++ b/src/nuclai/models/train_lr_rf.py @@ -88,7 +88,7 @@ def main(): grid_lr = ParameterGrid(param_grid_lr) for params in grid_lr: try: - lr = LogisticRegression(max_iter=10000, **params) + lr = LogisticRegression(max_iter=10000, n_jobs=-1, **params) lr.fit(data_train, labels_train) # Save the model penalty = params["penalty"] @@ -110,7 +110,7 @@ def main(): grid_rf = ParameterGrid(param_grid_rf) for params in grid_rf: try: - rf = RandomForestClassifier(**params) + rf = RandomForestClassifier(n_jobs=-1, **params) rf.fit(data_train, labels_train) # Save the model n_estimators = params["n_estimators"]