Saving samples plot as png at end of training

src/trainers/trainer.py

@@ -198,6 +198,7 @@ class Trainer:
 
             if task:
                 self.finish_training(task=task)
+                self.debug_plots(task, False, test_loader, test_samples, -1)
                 # task.close()
         except Exception:
             if task:
@@ -342,21 +343,8 @@ class Trainer:
                 features[:96], target, predictions, show_legend=(0 == 0)
             )
 
-            task.get_logger().report_matplotlib_figure(
-                title="Training" if train else "Testing",
-                series=f"Sample {actual_idx}",
-                iteration=epoch,
-                figure=fig,
-            )
-
-            task.get_logger().report_matplotlib_figure(
-                title="Training Samples" if train else "Testing Samples",
-                series=f"Sample {actual_idx} samples",
-                iteration=epoch,
-                figure=fig2,
-                report_interactive=False,
-            )
-
+            
+                )
             plt.close()
 
     def debug_scatter_plot(self, task, train: bool, samples, epoch):