From aaee16a632cbd552d33c9b797fa79df952a6bca3 Mon Sep 17 00:00:00 2001
From: "Daniel.Frisinghelli" <daniel.frisinghelli@eurac.edu>
Date: Wed, 3 Mar 2021 14:17:34 +0100
Subject: [PATCH] Compute confusion matrix on GPU, if available.
---
pysegcnn/core/trainer.py | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index a5d3506..dac20d2 100644
--- a/pysegcnn/core/trainer.py
+++ b/pysegcnn/core/trainer.py
@@ -2735,11 +2735,11 @@ class NetworkInference(BaseConfig):
# 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=np.asarray(list(self.use_labels.keys())))
+ conf_mat = confusion_matrix(y_true, y_pred, labels=labels_gpu)
# add confusion matrix to model output
output['cm'] = conf_mat.cpu().numpy()
--
GitLab