diff --git a/qim3d/utils/data.py b/qim3d/utils/data.py
index e50a3b87b6c2c3bbcadbbf2baf2e45fb5b3ca547..50f615541517c559887c33b5170fdc2879968391 100644
--- a/qim3d/utils/data.py
+++ b/qim3d/utils/data.py
@@ -145,7 +145,8 @@ class Dataset(torch.utils.data.Dataset):
             NotImplementedError: in Case 3, if a folder is found among the list of files.
             NotImplementedError: If the data structure does not fall into one of the three cases.
         '''
-        
+
+        # TODO: make sure user can specify extra mask names (add in prepare_datasets()?)
         target_folder_names = ['mask','label','target']
 
         # Case 1
@@ -186,6 +187,7 @@ class Dataset(torch.utils.data.Dataset):
             self.sample_targets = [target for target in sorted(Path(self.root_path,targets).iterdir())]
 
         # Case 3
+        # TODO: Make sure the split keeps the data from the same patient in either train, val or predict.
         elif self.n_folders > 2:
             sample_images = []
             sample_targets = []