[ ] Consider the generation of output using several hyperparam combinations for learning rate and batch size / n_epochs. Please include those in the fit_model wrapper as in the config.yml, to test several fitting scenarios and get conclusion on the type of training we need to perform across more datasets.