Formalize inputs to dataloader

[salomaestro]

I'm almost done with the model and dataset, but I see we need to formalize what inputs go to the different functions, load_model/data.

I've already pushed a few to the main branch, but I'll list them here, and we can decide what inputs we would like to the models/datasets:

Dataset

load_data(
    dataset_name: str,
    train: bool,
    data_path: pathlib.Path | str,
    download: bool,
)

Model

model = load_model(
    model_name: str,
    in_channels: int,  # So we can input both grayscale and RGB images
    num_classes: int,
)

Any other suggestions?

[hzavadil98]

At this point each model would be specifically tailored for the images from the connected dataset, why don't we pass the whole image.shape to the load_model so that we can make the models flexible to process all the implemented datasets? 🤔 @salomaestro

[salomaestro]

@hzavadil98 Good idea. If you have the time, could you open a pull-request/draft pull-request with your proposal?

[hzavadil98]

I'm on it