From 5b2c1ca58f0e524798d4b05f5813e292d47690d5 Mon Sep 17 00:00:00 2001 From: Adam Narozniak Date: Wed, 17 Jul 2024 12:59:48 +0200 Subject: [PATCH] Update tests for multiple partitioners --- datasets/flwr_datasets/federated_dataset_test.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/datasets/flwr_datasets/federated_dataset_test.py b/datasets/flwr_datasets/federated_dataset_test.py index f65aa6346f3a..e821cc359f86 100644 --- a/datasets/flwr_datasets/federated_dataset_test.py +++ b/datasets/flwr_datasets/federated_dataset_test.py @@ -144,9 +144,12 @@ def test_multiple_partitioners(self) -> None: dataset_test_partition0 = dataset_fds.load_partition(0, self.test_split) dataset = datasets.load_dataset(self.dataset_name) + expected_len = len(dataset[self.test_split]) // num_test_partitions + mod = len(dataset[self.test_split]) % num_test_partitions + expected_len += 1 if 0 < mod else 0 self.assertEqual( len(dataset_test_partition0), - len(dataset[self.test_split]) // num_test_partitions, + expected_len ) def test_no_need_for_split_keyword_if_one_partitioner(self) -> None: