Fix iterating over dataset keys in prediction stage

src/ptbench/data/datamodule.py

@@ -727,7 +727,7 @@ class CachingDataModule(lightning.LightningDataModule):
             self._setup_dataset("test")
 
         elif stage == "predict":
-            for k in self.database_split.keys():
+            for k in self.database_split:
                 self._setup_dataset(k)
 
     def teardown(self, stage: str) -> None: