diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index 12cdafbc822843973eeae49db4d93fce1e6572ed..a5d3506e3dc4a953daf7877aa272f09f5fe50ae1 100644
--- a/pysegcnn/core/trainer.py
+++ b/pysegcnn/core/trainer.py
@@ -2731,12 +2731,18 @@ 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)
+
+                # compute confusion matrix
                 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
+                output['cm'] = conf_mat.cpu().numpy()
 
                 # plot confusion matrix
                 plot_confusion_matrix(