fix inline error

This commit is contained in:
Falko Victor Habel 2025-01-26 23:04:49 +01:00
parent 7c4aef0978
commit 8a809269e5
1 changed files with 3 additions and 2 deletions

View File

@ -134,17 +134,18 @@ class AIIADataLoader:
self.items = [] self.items = []
for idx in range(len(dataset)): for idx in range(len(dataset)):
item = self.loader.get_item(idx) item = self.loader.get_item(idx)
if item is not None: if item is not None: # Only add valid items
if self.pretraining: if self.pretraining:
img = item[0] if isinstance(item, tuple) else item img = item[0] if isinstance(item, tuple) else item
self.items.append((img, 'denoise', img)) self.items.append((img, 'denoise', img))
self.items.append((img, 'rotate', 0)) self.items.append((img, 'rotate', 0))
else: else:
self.items.append(item) self.items.append(item)
if not self.items: if not self.items:
raise ValueError("No valid items were loaded from the dataset") raise ValueError("No valid items were loaded from the dataset")
train_indices, val_indices = self._split_data() train_indices, val_indices = self._split_data()
self.train_dataset = self._create_subset(train_indices) self.train_dataset = self._create_subset(train_indices)