call on_train_end when trial is pruned (#16536)

483a9450 · Ferdinand Schlatt · GitHub · 9de70f21 · 483a9450
Unverified Commit 483a9450 authored Apr 01, 2022 by Ferdinand Schlatt Committed by GitHub Apr 01, 2022
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

src/transformers/trainer.py src/transformers/trainer.py +1 -0

No files found.
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -991,6 +991,7 @@ class Trainer:

            trial.report(self.objective, epoch)
            if trial.should_prune():
+                self.callback_handler.on_train_end(self.args, self.state, self.control)
                raise optuna.TrialPruned()
        elif self.hp_search_backend == HPSearchBackend.RAY:
            from ray import tune