diff --git a/datasets/flwr_datasets/federated_dataset.py b/datasets/flwr_datasets/federated_dataset.py
index df92b79b90e4..069f2d5071b9 100644
--- a/datasets/flwr_datasets/federated_dataset.py
+++ b/datasets/flwr_datasets/federated_dataset.py
@@ -104,8 +104,8 @@ def __init__(
         # Indicate if the dataset is prepared for `load_partition` or `load_split`
         self._dataset_prepared: bool = False
         self._event = {
-            "load_partition": {split: False for split in self._partitioners.keys()},
-            "load_split": {split: False for split in self._partitioners.keys()},
+            "load_partition": {split: False for split in self._partitioners},
+            "load_split": {split: False for split in self._partitioners},
         }
 
     def load_partition(