diff --git a/src/aiia/data/DataLoader.py b/src/aiia/data/DataLoader.py
index 567954f..cbb22a6 100644
--- a/src/aiia/data/DataLoader.py
+++ b/src/aiia/data/DataLoader.py
@@ -134,17 +134,18 @@ class AIIADataLoader:
         self.items = []
         for idx in range(len(dataset)):
             item = self.loader.get_item(idx)
-            if item is not None:
+            if item is not None:  # Only add valid items
                 if self.pretraining:
                     img = item[0] if isinstance(item, tuple) else item
                     self.items.append((img, 'denoise', img))
                     self.items.append((img, 'rotate', 0))
-            else:
+                else:
                     self.items.append(item)
         
         if not self.items:
             raise ValueError("No valid items were loaded from the dataset")
 
+
         train_indices, val_indices = self._split_data()
         
         self.train_dataset = self._create_subset(train_indices)