diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index be12dd3142bf7bc717d7246ffd3debc73fcece86..d09a6c31ee3c18739bd7778a0122422ced8e5d3e 100644
--- a/pysegcnn/core/trainer.py
+++ b/pysegcnn/core/trainer.py
@@ -2607,7 +2607,7 @@ class NetworkInference(BaseConfig):
                 domain=self.domain)
 
             # load the source dataset the model was trained on
-            self.src_ds = self.load_dataset(state, test=None)
+            self.src_ds = self.load_dataset(state, test=None).dataset
 
             # load the pretrained model
             model, _ = Network.load_pretrained_model(state)