Cleaning up mlp_train code