Add numpy import to docstring

adap · Jul 17, 2024 · 2501830 · 2501830
1 parent 06ae94d
commit 2501830
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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