Fix extra channel in a saved image

config_base.yaml

@@ -19,7 +19,7 @@ step: 2
 
 # Training related
 batch_size: 32
-num_epochs: 20
+num_epochs: 40
 resume: ''
 load_optim: True
 anew: False

src/core/trainers.py

@@ -267,6 +267,6 @@ class CDTrainer(Trainer):
                 self.logger.dump(desc)
                     
                 if store:
-                    self.save_image(name[0], CM, epoch)
+                    self.save_image(name[0], CM.squeeze(-1), epoch)
 
         return self.metrics[0].avg if len(self.metrics) > 0 else max(1.0 - losses.avg, self._init_max_acc)
\ No newline at end of file