diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index 8e252022a7d6ee90b4507d2512bb094b984def45..502551e67254d7608775a09ff6d861fb2c70d732 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().flatten(axis=0)
-            labels = labels.cpu().numpy().flatten(axis=0)
-            prdctn = prdctn.cpu().numpy().flatten(axis=0)
+            inputs = inputs.cpu().numpy().squeeze(axis=0)
+            labels = labels.cpu().numpy().squeeze(axis=0)
+            prdctn = prdctn.cpu().numpy().squeeze(axis=0)
 
             # check whether to reconstruct the scenes of a dataset
             if self.predict_scene: