def METHOD_NAME(self, _: str, train_num_partitions: int) -> None:
    """Test if the partition size is correct based on the number of partitions."""
    dataset_fds = FederatedDataset(
        dataset=self.dataset_name, partitioners={"train": train_num_partitions}
    )
    dataset_partition0 = dataset_fds.load_partition(0, "train")
    dataset = datasets.load_dataset(self.dataset_name)
    self.assertEqual(
        len(dataset_partition0), len(dataset["train"]) // train_num_partitions
    )