Skip to content

Commit 7d92009

Browse files
authored
Added additional kwarg for successful running of optuna hyperparameter search (#31924)
Update integration_utils.py Added additional kwarg
1 parent 6370062 commit 7d92009

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

src/transformers/integrations/integration_utils.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -252,10 +252,11 @@ def _objective(trial, checkpoint_dir=None):
252252

253253
timeout = kwargs.pop("timeout", None)
254254
n_jobs = kwargs.pop("n_jobs", 1)
255+
gc_after_trial = kwargs.pop("gc_after_trial", False)
255256
directions = direction if isinstance(direction, list) else None
256257
direction = None if directions is not None else direction
257258
study = optuna.create_study(direction=direction, directions=directions, **kwargs)
258-
study.optimize(_objective, n_trials=n_trials, timeout=timeout, n_jobs=n_jobs)
259+
study.optimize(_objective, n_trials=n_trials, timeout=timeout, n_jobs=n_jobs, gc_after_trial=gc_after_trial)
259260
if not study._is_multi_objective():
260261
best_trial = study.best_trial
261262
return BestRun(str(best_trial.number), best_trial.value, best_trial.params)

0 commit comments

Comments
 (0)