Make sure on_save of extensions is called on TrainCV save