From 57b8b57d46755329ae2520095c608ad337d37dda Mon Sep 17 00:00:00 2001
From: dcarron <daniel.carron@idiap.ch>
Date: Mon, 24 Jul 2023 11:12:42 +0200
Subject: [PATCH] Fix iterating over dataset keys in prediction stage
---
src/ptbench/data/datamodule.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/src/ptbench/data/datamodule.py b/src/ptbench/data/datamodule.py
index 440c1ea4..0fb4edcf 100644
--- a/src/ptbench/data/datamodule.py
+++ b/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:
--
GitLab