From 5427d7e9842d68a5e4e25aabf5c3e0ae30a8e26d Mon Sep 17 00:00:00 2001
From: Bobholamovic <bob1998425@hotmail.com>
Date: Mon, 9 Dec 2019 09:17:38 +0800
Subject: [PATCH] Fix extra channel in a saved image
---
config_base.yaml | 2 +-
src/core/trainers.py | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/config_base.yaml b/config_base.yaml
index 23154e1..f581c71 100644
--- a/config_base.yaml
+++ b/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
diff --git a/src/core/trainers.py b/src/core/trainers.py
index ddefda4..2c5b139 100644
--- a/src/core/trainers.py
+++ b/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
--
GitLab