diff --git a/kraken/lib/pretrain/model.py b/kraken/lib/pretrain/model.py
index 97b220751..eee832b6d 100644
--- a/kraken/lib/pretrain/model.py
+++ b/kraken/lib/pretrain/model.py
@@ -203,11 +203,11 @@ def __init__(self,
             self.train_set, self.val_set = random_split(train_set, (train_len, val_len))
 
         if format_type == 'binary':
-            legacy_train_status = train_set.legacy_polygons_status
-            if val_set and val_set.legacy_polygons_status != legacy_train_status:
+            legacy_train_status = self.train_set.dataset.legacy_polygons_status
+            if self.val_set.dataset.legacy_polygons_status != legacy_train_status:
                 logger.warning('Train and validation set have different legacy '
                                f'polygon status: {legacy_train_status} and '
-                               f'{val_set.legacy_polygons_status}. Train set '
+                               f'{self.val_set.dataset.legacy_polygons_status}. Train set '
                                'status prevails.')
             if legacy_train_status == "mixed":
                 logger.warning('Mixed legacy polygon status in training dataset. Consider recompilation.')
diff --git a/kraken/lib/train.py b/kraken/lib/train.py
index e7506658c..dca4b6def 100644
--- a/kraken/lib/train.py
+++ b/kraken/lib/train.py
@@ -394,11 +394,11 @@ def __init__(self,
                              'command. Please add valid XML, line, or binary data.')
 
         if format_type == 'binary':
-            legacy_train_status = train_set.legacy_polygons_status
-            if val_set and val_set.legacy_polygons_status != legacy_train_status:
+            legacy_train_status = self.train_set.dataset.legacy_polygons_status
+            if self.val_set.dataset.legacy_polygons_status != legacy_train_status:
                 logger.warning('Train and validation set have different legacy '
                                f'polygon status: {legacy_train_status} and '
-                               f'{val_set.legacy_polygons_status}. Train set '
+                               f'{self.val_set.dataset.legacy_polygons_status}. Train set '
                                'status prevails.')
             if legacy_train_status == "mixed":
                 logger.warning('Mixed legacy polygon status in training dataset. Consider recompilation.')