diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index c3066ad4717d46917343517ab25890709431dfaf..f3089ff04f3d6ef9d1c9714e8f425a077fc9e96a 100644
--- a/pysegcnn/core/trainer.py
+++ b/pysegcnn/core/trainer.py
@@ -2513,6 +2513,11 @@ class NetworkInference(BaseConfig):
             progress = 'Sample: {:d}/{:d}'.format(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()
+
             # check whether to reconstruct the scene
             if self.dataloader.batch_size > 1:
 
@@ -2643,9 +2648,9 @@ class NetworkInference(BaseConfig):
             if self.cm:
 
                 # merge predictions for all samples
-                y_true = np.asarray([v['y_true'].numpy().flatten() for _, v
+                y_true = np.asarray([v['y_true'].flatten() for _, v
                                      in output.items()]).flatten()
-                y_pred = np.asarray([v['y_pred'].numpy().flatten() for _, v
+                y_pred = np.asarray([v['y_pred'].flatten() for _, v
                                      in output.items()]).flatten()
 
                 # calculate confusion matrix
diff --git a/pysegcnn/core/utils.py b/pysegcnn/core/utils.py
index 24d905b09e102c60b3306db4115ca0bcdbc79dba..609deda62820411e380d612a9ac48bb9d351eb65 100644
--- a/pysegcnn/core/utils.py
+++ b/pysegcnn/core/utils.py
@@ -745,7 +745,7 @@ def reconstruct_scene(tiles):
 
     Parameters
     ----------
-    tiles : :py:class:`torch.Tensor` or :py:class:`numpy.ndarray`
+    tiles : :py:class:`numpy.ndarray`
         The tiled image, shape: `(tiles, bands, tile_size, tile_size)` or
         `(tiles, tile_size, tile_size)`.
 
@@ -755,32 +755,26 @@ def reconstruct_scene(tiles):
         The reconstructed image, shape: `(bands, height, width)`.
 
     """
-    # check if tensor is on gpu and convert to numpy array
-    if isinstance(tiles, torch.Tensor):
-        tiles_cpu = np.asarray(tiles.cpu())
-    else:
-        tiles_cpu = np.asarray(tiles)
-
     # check the dimensions of the input array
-    if tiles_cpu.ndim > 3:
-        nbands = tiles_cpu.shape[1]
-        tile_size = tiles_cpu.shape[2]
+    if tiles.ndim > 3:
+        nbands = tiles.shape[1]
+        tile_size = tiles.shape[2]
     else:
         nbands = 1
-        tile_size = tiles_cpu.shape[1]
+        tile_size = tiles.shape[1]
 
     # calculate image size
-    img_size = 2 * (int(np.sqrt(tiles_cpu.shape[0]) * tile_size),)
+    img_size = 2 * (int(np.sqrt(tiles.shape[0]) * tile_size),)
 
     # calculate the topleft corners of the tiles
     topleft = tile_topleft_corner(img_size, tile_size)
 
     # iterate over the tiles
     scene = np.zeros(shape=(nbands,) + img_size)
-    for t in range(tiles_cpu.shape[0]):
+    for t in range(tiles.shape[0]):
         scene[...,
               topleft[t][0]: topleft[t][0] + tile_size,
-              topleft[t][1]: topleft[t][1] + tile_size] = tiles_cpu[t, ...]
+              topleft[t][1]: topleft[t][1] + tile_size] = tiles[t, ...]
 
     return scene.squeeze()