As part of understanding the script enough to get a sweep running from my machine I was struggling because of how much code there was in the train function, and how model configuration was split between ICLCOnfig defaults and explicit values in get_config function.
I extracted all evaluations functionality from training function into a new Evaluator class. An Evaluator stores the fixed evaluation batches and provides a method that takes a model and produces the metrics dictionary
I have also attempted to inline losses_fn so that all of the evaluation code fits in the one function
I moved all model and data set configuration into get_config rather than the default values of the ICLConfig class (which now has no default values, so all values should be provided explicitly within get_config)
A small number of other miscellaneous local changes were made (see diff and comments on GitHub PR).
As part of understanding the script enough to get a sweep running from my machine I was struggling because of how much code there was in the train function, and how model configuration was split between ICLCOnfig defaults and explicit values in get_config function.
losses_fn
so that all of the evaluation code fits in the one functionA small number of other miscellaneous local changes were made (see diff and comments on GitHub PR).