diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index 2640ca953ad6796bd58ff0e5096007ae6afab2ae..689f9c5306929721bf3e2c59f582ee5de4c0694b 100644
--- a/pysegcnn/core/trainer.py
+++ b/pysegcnn/core/trainer.py
@@ -2593,9 +2593,10 @@ class NetworkInference(BaseConfig):
                     if self.plot:
 
                         # title for prediction
-                        title = ''.join([(v[0] + str(k)) for k, v in
-                                         self.src_ds.sensor.band_dict().items()
-                                         if v in self.bands])
+                        title = ''.join([
+                            (v[0] + str(k)) for k, v in
+                            self.src_ds.dataset.sensor.band_dict().items()
+                            if v in self.bands])
 
                         # plot inputs, ground truth and model predictions
                         fig = plot_sample(inputs.clip(0, 1),