Update __init__.py

ultralytics · Oct 18, 2021 · 01e1be8 · 01e1be8
1 parent 1db6b9d
commit 01e1be8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py
@@ -131,7 +131,7 @@ def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
             if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1:
                 self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
 
-    def on_train_end(self, last, best, plots, epoch):
+    def on_train_end(self, last, best, plots, epoch, results):
         # Callback runs on training end
         if plots:
             plot_results(file=self.save_dir / 'results.csv')  # save results.png