fixing create_federated_dataloaders to work with iid option

[cozek]

Issue

None

Description

In dataset.py of the leaf baseline, the idd sampling type throws an error. This is because it doesn't take the n_clients parameter which is required for iid. I added that parameter.

trainloaders, valloaders, testloaders = create_federated_dataloaders(
    sampling_type = 'iid',
    dataset_fraction = 1,
    batch_size = 32,
    train_fraction = 0.8,
    validation_fraction = 0.1,
    test_fraction = 0.1,
    random_seed = 42,
    n_clients=10
)

Error: "n_clients can not be None for idd training"

Related issues/PRs

None

Proposal

Add the n_clients parameter to the function and add an if else statement for iid and non-idd.

Explanation

I added the n_clients parameter and made the necessary changes in the function.

trainloaders, valloaders, testloaders = create_federated_dataloaders(
    sampling_type = 'iid',
    dataset_fraction = 1,
    batch_size = 32,
    train_fraction = 0.8,
    validation_fraction = 0.1,
    test_fraction = 0.1,
    random_seed = 42,
)

Checklist

• [X] Implement proposed change
• [ ] Write tests
• [ ] Update documentation
• [ ] Update changelog
• [ ] Make CI checks pass
• [ ] Ping maintainers on Slack (channel #contributions)

Any other comments?

Not sure how to write the tests for this. idd test case is missing from dataset_test.py