Fixed losses and multiple other things

src/trainers/trainer.py

@@ -313,6 +313,7 @@ class Trainer:
     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(
@@ -341,6 +342,7 @@ class Trainer:
             for trace in sub_fig.data:
                 fig.add_trace(trace, row=row, col=col)
 
+
             # 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}")