* Fix training * Add test * Fix failing tests Co-authored-by: Niels Rogge <nielsrogge@Nielss-MacBook-Pro.local>