Add support for multi-task models

[ivyleavedtoadflax]

The original Rodrigues et al. (2018) model was a multitask model that combined three tasks in a single training cycle. So far this package has been built to accomodate separate models. There are a few adaptations that must be made to allow the code to be re-used for a multi-task model.