diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index 7ebcd1e5b778dff6e2392c56c1f3a5c2aa20c9cb..8e252022a7d6ee90b4507d2512bb094b984def45 100644
--- a/pysegcnn/core/trainer.py
+++ b/pysegcnn/core/trainer.py
@@ -2538,9 +2538,9 @@ class NetworkInference(BaseConfig):
                 batch + 1, len(self.dataloader))
 
             # check if tensor is on gpu and convert to numpy array
-            inputs = inputs.cpu().numpy()
-            labels = labels.cpu().numpy()
-            prdctn = prdctn.cpu().numpy()
+            inputs = inputs.cpu().numpy().flatten(axis=0)
+            labels = labels.cpu().numpy().flatten(axis=0)
+            prdctn = prdctn.cpu().numpy().flatten(axis=0)
 
             # check whether to reconstruct the scenes of a dataset
             if self.predict_scene: