Updated some stuff

src/trainers/trainer.py

@@ -7,6 +7,7 @@ import numpy as np
 from plotly.subplots import make_subplots
 from clearml.config import running_remotely
 from torchinfo import summary
+import matplotlib.pyplot as plt
 
 
 class Trainer:
@@ -329,18 +330,7 @@ class Trainer:
         return fig
 
     def debug_plots(self, task, train: bool, data_loader, sample_indices, epoch):
-        num_samples = len(sample_indices)
-        rows = num_samples  # One row per sample since we only want one column
-
-        cols = 1
-
-        fig = make_subplots(
-            rows=rows,
-            cols=cols,
-            subplot_titles=[f"Sample {i+1}" for i in range(num_samples)],
-        )
-
-        for i, idx in enumerate(sample_indices):
+        for actual_idx, idx in sample_indices.items():
             features, target, _ = data_loader.dataset[idx]
 
             features = features.to(self.device)
@@ -350,30 +340,26 @@ class Trainer:
             with torch.no_grad():
                 predictions = self.model(features).cpu()
 
-            sub_fig = self.get_plot(
-                features[:96], target, predictions, show_legend=(i == 0)
+            fig, fig2 = self.get_plot(
+                features[:96], target, predictions, show_legend=(0 == 0)
             )
 
-            row = i + 1
-            col = 1
+            task.get_logger().report_matplotlib_figure(
+                title="Training" if train else "Testing",
+                series=f"Sample {actual_idx}",
+                iteration=epoch,
+                figure=fig,
+            )
 
-            for trace in sub_fig.data:
-                fig.add_trace(trace, row=row, col=col)
+            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,
+            )
 
-            # loss = self.criterion(predictions.to(self.device), target.squeeze(-1).to(self.device)).item()
-
-            # fig['layout']['annotations'][i].update(text=f"{loss.__class__.__name__}: {loss:.6f}")
-
-        # y axis same for all plots
-        # fig.update_yaxes(range=[-1, 1], col=1)
-
-        fig.update_layout(height=1000 * rows)
-        task.get_logger().report_plotly(
-            title=f"{'Training' if train else 'Test'} Samples",
-            series="full_day",
-            iteration=epoch,
-            figure=fig,
-        )
+            plt.close()
 
     def debug_scatter_plot(self, task, train: bool, samples, epoch):
         X, y = samples