diff --git a/datasets/flwr_datasets/partitioner/image_semantic_partitioner.py b/datasets/flwr_datasets/partitioner/image_semantic_partitioner.py
index 20ba9f9d2cc7..2bc29f188494 100644
--- a/datasets/flwr_datasets/partitioner/image_semantic_partitioner.py
+++ b/datasets/flwr_datasets/partitioner/image_semantic_partitioner.py
@@ -96,6 +96,10 @@ class ImageSemanticPartitioner(Partitioner):
         Seed used for PCA dimensionality reduction. If not set, sample_seed is used.
     gmm_seed: Optional[int]
         Seed used for GMM clustering. If not set, sample_seed is used.
+    image_column_name: Optional[str]
+        The name of the image column in the dataset. If not set, the first image column
+        is used.
+
 
     Examples
     --------
@@ -498,4 +502,4 @@ def _check_variable_validation(self) -> None:
         if not isinstance(self._pca_seed, int):
             raise TypeError("The pca seed needs to be an integer.")
         if not isinstance(self._gmm_seed, int):
-            raise TypeError("The gmm seed needs to be an integer.")
\ No newline at end of file
+            raise TypeError("The gmm seed needs to be an integer.")