[BugFix] Fix Test wandb callback bug (PaddlePaddle#8056)

zhangbo9674 · Mar 6, 2024 · 32d8b3f · 32d8b3f
1 parent 092c845
commit 32d8b3f
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/trainer/test_trainer_visualization.py b/tests/trainer/test_trainer_visualization.py
@@ -60,6 +60,7 @@ def test_wandbcallback(self):
                 self.assertEqual(wandbcallback._wandb.run.summary["train/learning_rate"], log["learning_rate"])
                 self.assertEqual(wandbcallback._wandb.run.summary["train/global_step"], log["global_step"])
         wandbcallback.on_train_end(args, state, control, model=model)
+        wandbcallback._wandb.finish()
         os.environ.pop("WANDB_LOG_MODEL", None)
         os.environ.pop("WANDB_MODE", None)
         shutil.rmtree(output_dir)