From 06199357afe341dc2db7d08bde0b097c971729b3 Mon Sep 17 00:00:00 2001
From: Liunaijiaaa <liunaijiaaa@yeah.net>
Date: Sat, 13 Nov 2021 16:41:05 +0800
Subject: [PATCH] Solved the problem that wandb could not display masks_pred in
 the multi-category segmentation problem.If we do not modify it, we may
 encounter the following error message:mode, rawmode =
 _fromarray_typemap[typekey] KeyError: ((1, 1, 15), '|u1').ps:15 is the number
 of categories.

Former-commit-id: ff0c2ffa33c945a491e85ef44e89891e9a7c435d
---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index 146373a..cae9a86 100644
--- a/train.py
+++ b/train.py
@@ -130,7 +130,7 @@ def train_net(net,
                             'images': wandb.Image(images[0].cpu()),
                             'masks': {
                                 'true': wandb.Image(true_masks[0].float().cpu()),
-                                'pred': wandb.Image(torch.softmax(masks_pred, dim=1)[0].float().cpu()),
+                                'pred': wandb.Image(torch.softmax(masks_pred, dim=1).argmax(dim=1)[0].float().cpu()),
                             },
                             'step': global_step,
                             'epoch': epoch,