Fix display
Train last layer correctly