* allow custom trainer_cls to be defined as a module reference in the YAML * address PR feedback and add test * add tests