Skip to content

Commit 4dcd1ce

Browse files
committed
Add num-threads passing for xgboost benchmark
1 parent 3e9be64 commit 4dcd1ce

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

xgboost/gbt.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,9 @@ def convert_xgb_predictions(y_pred, objective):
9494
'seed': params.seed
9595
}
9696

97+
if params.threads != -1:
98+
xgb_params.update({'nthread': params.threads})
99+
97100
columns = ('batch', 'arch', 'prefix', 'function', 'threads', 'dtype', 'size',
98101
'num_trees')
99102

0 commit comments

Comments
 (0)