Only log total_flos at the end of training (#7981)

* Only log total_flos at the end of training * Fix test
2020-10-22 14:26:55 -04:00
parent ff65beafa3
commit 06fc3954a1
3 changed files with 6 additions and 10 deletions
--- a/tests/test_trainer_callback.py
+++ b/tests/test_trainer_callback.py
@@ -125,7 +125,7 @@ class TrainerCallbackTest(unittest.TestCase):
            expected_events.append("on_epoch_end")
            if trainer.args.evaluation_strategy == EvaluationStrategy.EPOCH:
                expected_events += evaluation_events.copy()
-        expected_events.append("on_train_end")
+        expected_events += ["on_log", "on_train_end"]
        return expected_events

    def test_init_callback(self):