Squeeze is required for multi-dimensional arrays.

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: