diff --git a/src/aiia/data/DataLoader.py b/src/aiia/data/DataLoader.py
index 1813749..98b882b 100644
--- a/src/aiia/data/DataLoader.py
+++ b/src/aiia/data/DataLoader.py
@@ -165,7 +165,8 @@ class AIIADataLoader:
         if len(self.items) == 0:
             return [], []
             
-        tasks = [item[1] if len(item) > 1 and hasattr(item, '__getitem__') else None for item in self.items]
+        tasks = [item[1] for item in self.items if len(item) > 1 and hasattr(item, '__getitem__') and item[1] is not None]
+
         unique_tasks = list(set(tasks)) if tasks.count(None) < len(tasks) else []
         
         train_indices = []