Wrote more of the training code