diff --git a/datasets/flwr_datasets/partitioner/distribution_partitioner.py b/datasets/flwr_datasets/partitioner/distribution_partitioner.py
index da540f236510..e01740091b34 100644
--- a/datasets/flwr_datasets/partitioner/distribution_partitioner.py
+++ b/datasets/flwr_datasets/partitioner/distribution_partitioner.py
@@ -92,6 +92,7 @@ class label distribution. `float` values are rounded to the nearest `int`.
     >>> from flwr_datasets import FederatedDataset
     >>> from flwr_datasets.partitioner import DistributionPartitioner
     >>> from pprint import pprint
+    >>> import numpy as np
     >>>
     >>> num_clients = 1_000
     >>> num_unique_labels_per_client = 2