diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index 206e780c149c9a19a0875b746091d5700dc4db6b..a32fa26a406a573ea31ed5a9725b239e5549b665 100644
--- a/pysegcnn/core/trainer.py
+++ b/pysegcnn/core/trainer.py
@@ -2532,10 +2532,6 @@ class NetworkInference(BaseConfig):
                 prdctn = F.softmax(
                     model(inputs), dim=1).argmax(dim=1).squeeze()
 
-            # check whether the source and target domain labels differ
-            if self.apply_label_map:
-                prdctn = self.map_to_target(prdctn)
-
             # progress string to log
             progress = 'Sample: {:d}/{:d}'.format(
                 batch + 1, len(self.dataloader))
@@ -2545,6 +2541,10 @@ class NetworkInference(BaseConfig):
             labels = labels.cpu().numpy().squeeze()
             prdctn = prdctn.cpu().numpy().squeeze()
 
+            # check whether the source and target domain labels differ
+            if self.apply_label_map:
+                prdctn = self.map_to_target(prdctn)
+
             # check whether to reconstruct the scenes of a dataset
             if self.predict_scene: