diff --git a/pyproject.toml b/pyproject.toml
index 6b381da2be1687c18da9095aada87ec12745777e..e43899d2dbad33e10c8b4eb67a333f12547aca55 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -41,7 +41,7 @@ dependencies = [
    "pillow",
    "torch>=1.8",
    "torchvision>=0.10",
-   "lightning",
+   "lightning>=2.0.3",
    "tensorboard",
 ]
 
@@ -117,8 +117,8 @@ 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.configs.datasets.shenzhen.default"
-shenzhen_rgb = "ptbench.configs.datasets.shenzhen.rgb"
+shenzhen = "ptbench.data.shenzhen.default"
+shenzhen_rgb = "ptbench.data.shenzhen.rgb"
 shenzhen_f0 = "ptbench.data.shenzhen.fold_0"
 shenzhen_f1 = "ptbench.data.shenzhen.fold_1"
 shenzhen_f2 = "ptbench.data.shenzhen.fold_2"