vector_norm(dim=1) dropped the channel axis, so dividing the (B,C,H,W)
latent by the (B,H,W) norm broadcast the batch axis onto the channels.
With batch_size != channels this raised a RuntimeError, and when they
happened to be equal it silently mixed the per-sample normalization.
Keep the channel dimension with keepdim=True so each latent is
normalized independently.
