Skip to content

GitLab

Unverified Commit 483a9450 authored by Ferdinand Schlatt's avatar Ferdinand Schlatt Committed by GitHub
Browse files

call on_train_end when trial is pruned (#16536)

parent 9de70f21
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
src/transformers/trainer.py
View file @ 483a9450
...@@ -991,6 +991,7 @@ class Trainer: ...@@ -991,6 +991,7 @@ class Trainer:
trial.report(self.objective, epoch) trial.report(self.objective, epoch)
if trial.should_prune(): if trial.should_prune():
self.callback_handler.on_train_end(self.args, self.state, self.control)
raise optuna.TrialPruned() raise optuna.TrialPruned()
elif self.hp_search_backend == HPSearchBackend.RAY: elif self.hp_search_backend == HPSearchBackend.RAY:
from ray import tune from ray import tune
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment