Change the check to allow using only all columns

adap · Dec 13, 2024 · ac3bc68 · ac3bc68
1 parent 8634269
commit ac3bc68
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py b/datasets/flwr_datasets/partitioner/vertical_size_partitioner.py
@@ -259,9 +259,9 @@ def _validate_parameters_while_partitioning(
                 )
         num_columns = len(all_columns)
         if all(isinstance(size, int) for size in self._partition_sizes):
-            if sum(self._partition_sizes) > num_columns:
+            if sum(self._partition_sizes) != num_columns:
                 raise ValueError(
-                    "Sum of partition sizes cannot exceed the total number of columns."
+                    "Sum of partition sizes cannot differ from the total number of columns."
                 )
         else:
             pass