diff --git a/datasets/flwr_datasets/partitioner/natural_id_partitioner.py b/datasets/flwr_datasets/partitioner/natural_id_partitioner.py
index d258ebc32790..f6a44cd8e622 100644
--- a/datasets/flwr_datasets/partitioner/natural_id_partitioner.py
+++ b/datasets/flwr_datasets/partitioner/natural_id_partitioner.py
@@ -153,7 +153,7 @@ def partition_id_to_natural_id(self, value: Dict[int, str]) -> None:
         )
 
     def _check_supported_type_of_value_in_partition_by(self) -> None:
-        values = self.dataset[self._partition_by]
+        values = self.dataset[0][self._partition_by]
         values_np = np.array(values)
         dtype = values_np.dtype
         if not (