build_network now returns the pre-logits, so~we'll change