Improve train model