diff --git a/utils/datasets.py b/utils/datasets.py
index 4658dc524be0..c2859a148106 100755
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -1116,7 +1116,7 @@ def round_labels(labels):
     nc = data['nc']  # number of classes
     stats = {'nc': nc, 'names': data['names']}  # statistics dictionary
     for split in 'train', 'val', 'test':
-        if split not in data:
+        if data.get(split) is None:
             stats[split] = None  # i.e. no test set
             continue
         x = []