From aa7a0e91b175a9b2c64c3eca44dcafe77dc23d03 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Thu, 5 May 2022 21:25:13 -0700
Subject: [PATCH] Update `on_train_end` callback (#7716)

---
 train.py                  | 1 -
 utils/loggers/__init__.py | 1 +
 2 files changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index 86247c4df5f7..707651637a79 100644
--- a/train.py
+++ b/train.py
@@ -466,7 +466,6 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
                         callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
 
         callbacks.run('on_train_end', last, best, plots, epoch, results)
-        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
 
     torch.cuda.empty_cache()
     return results
diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py
index b33b04e46752..ef94615e9e57 100644
--- a/utils/loggers/__init__.py
+++ b/utils/loggers/__init__.py
@@ -164,6 +164,7 @@ def on_train_end(self, last, best, plots, epoch, results):
             plot_results(file=self.save_dir / 'results.csv')  # save results.png
         files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))]
         files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()]  # filter
+        self.logger.info(f"Results saved to {colorstr('bold', self.save_dir)}")
 
         if self.tb:
             for f in files: