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