From bf0b5ba3dc630d2aa5c113b44fdafbbb28512bba Mon Sep 17 00:00:00 2001
From: "Daniel.Frisinghelli" <daniel.frisinghelli@eurac.edu>
Date: Wed, 3 Mar 2021 14:24:22 +0100
Subject: [PATCH] Computing cm on GPU requires rewriting sklearn cm. Aborting.

---
 pysegcnn/core/trainer.py | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index dac20d2..12cdafb 100644
--- a/pysegcnn/core/trainer.py
+++ b/pysegcnn/core/trainer.py
@@ -2731,18 +2731,12 @@ class NetworkInference(BaseConfig):
 
                 # calculate confusion matrix
                 LOGGER.info('Computing confusion matrix ...')
-
-                # move predictions and labels to GPU if available
-                y_true = torch.Tensor(y_true).to(self.device)
-                y_pred = torch.Tensor(y_pred).to(self.device)
-                labels_gpu = torch.Tensor(
-                    np.asarray(list(self.use_labels.keys()))).to(self.device)
-
-                # compute confusion matrix
-                conf_mat = confusion_matrix(y_true, y_pred, labels=labels_gpu)
+                conf_mat = confusion_matrix(
+                    y_true, y_pred,
+                    labels=np.asarray(list(self.use_labels.keys())))
 
                 # add confusion matrix to model output
-                output['cm'] = conf_mat.cpu().numpy()
+                output['cm'] = conf_mat
 
                 # plot confusion matrix
                 plot_confusion_matrix(
-- 
GitLab