diff --git a/pyproject.toml b/pyproject.toml
index 7f3a5aaa75d46e79ab893648b30a49e3d9228923..3241d527546436d0feeeefa31bc4d8ed5a48495c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -77,17 +77,17 @@ alexnet_pre = "ptbench.configs.models.alexnet_pretrained"
 densenet = "ptbench.configs.models.densenet"
 densenet_pre = "ptbench.configs.models.densenet_pretrained"
 # montgomery dataset (and cross-validation folds)
-montgomery = "ptbench.data.montgomery.default"
-montgomery_f0 = "ptbench.data.montgomery.fold_0"
-montgomery_f1 = "ptbench.data.montgomery.fold_1"
-montgomery_f2 = "ptbench.data.montgomery.fold_2"
-montgomery_f3 = "ptbench.data.montgomery.fold_3"
-montgomery_f4 = "ptbench.data.montgomery.fold_4"
-montgomery_f5 = "ptbench.data.montgomery.fold_5"
-montgomery_f6 = "ptbench.data.montgomery.fold_6"
-montgomery_f7 = "ptbench.data.montgomery.fold_7"
-montgomery_f8 = "ptbench.data.montgomery.fold_8"
-montgomery_f9 = "ptbench.data.montgomery.fold_9"
+montgomery = "ptbench.data.montgomery.default:datamodule"
+montgomery_f0 = "ptbench.data.montgomery.fold_0:datamodule"
+montgomery_f1 = "ptbench.data.montgomery.fold_1:datamodule"
+montgomery_f2 = "ptbench.data.montgomery.fold_2:datamodule"
+montgomery_f3 = "ptbench.data.montgomery.fold_3:datamodule"
+montgomery_f4 = "ptbench.data.montgomery.fold_4:datamodule"
+montgomery_f5 = "ptbench.data.montgomery.fold_5:datamodule"
+montgomery_f6 = "ptbench.data.montgomery.fold_6:datamodule"
+montgomery_f7 = "ptbench.data.montgomery.fold_7:datamodule"
+montgomery_f8 = "ptbench.data.montgomery.fold_8:datamodule"
+montgomery_f9 = "ptbench.data.montgomery.fold_9:datamodule"
 # extended montgomery dataset (with radiological signs)
 montgomery_rs = "ptbench.configs.datasets.montgomery_RS.default"
 montgomery_rs_f0 = "ptbench.configs.datasets.montgomery_RS.fold_0"
@@ -101,17 +101,17 @@ montgomery_rs_f7 = "ptbench.configs.datasets.montgomery_RS.fold_7"
 montgomery_rs_f8 = "ptbench.configs.datasets.montgomery_RS.fold_8"
 montgomery_rs_f9 = "ptbench.configs.datasets.montgomery_RS.fold_9"
 # shenzhen dataset (and cross-validation folds)
-shenzhen = "ptbench.data.shenzhen.default"
-shenzhen_f0 = "ptbench.data.shenzhen.fold_0"
-shenzhen_f1 = "ptbench.data.shenzhen.fold_1"
-shenzhen_f2 = "ptbench.data.shenzhen.fold_2"
-shenzhen_f3 = "ptbench.data.shenzhen.fold_3"
-shenzhen_f4 = "ptbench.data.shenzhen.fold_4"
-shenzhen_f5 = "ptbench.data.shenzhen.fold_5"
-shenzhen_f6 = "ptbench.data.shenzhen.fold_6"
-shenzhen_f7 = "ptbench.data.shenzhen.fold_7"
-shenzhen_f8 = "ptbench.data.shenzhen.fold_8"
-shenzhen_f9 = "ptbench.data.shenzhen.fold_9"
+shenzhen = "ptbench.data.shenzhen.default:datamodule"
+shenzhen_f0 = "ptbench.data.shenzhen.fold_0:datamodule"
+shenzhen_f1 = "ptbench.data.shenzhen.fold_1:datamodule"
+shenzhen_f2 = "ptbench.data.shenzhen.fold_2:datamodule"
+shenzhen_f3 = "ptbench.data.shenzhen.fold_3:datamodule"
+shenzhen_f4 = "ptbench.data.shenzhen.fold_4:datamodule"
+shenzhen_f5 = "ptbench.data.shenzhen.fold_5:datamodule"
+shenzhen_f6 = "ptbench.data.shenzhen.fold_6:datamodule"
+shenzhen_f7 = "ptbench.data.shenzhen.fold_7:datamodule"
+shenzhen_f8 = "ptbench.data.shenzhen.fold_8:datamodule"
+shenzhen_f9 = "ptbench.data.shenzhen.fold_9:datamodule"
 # extended shenzhen dataset (with radiological signs)
 shenzhen_rs = "ptbench.data.shenzhen_RS.default"
 shenzhen_rs_f0 = "ptbench.configs.datasets.shenzhen_RS.fold_0"