bug fix model training