About Mapper in Latent Mapper, does PixelNorm normalize the right dim?

[Big-Brother-Pikachu]

Thanks for sharing this wonderful work! Here I am a little confused with the normalization layer: https://github.com/orpatashnik/StyleCLIP/blob/adc836d9077b241d92e9ce88da8cb6c01ada2999/mapper/latent_mappers.py#L16 In the StyleGAN2 mapper, as the input is always (Batch, latent_dim), the normalization layer is right here. https://github.com/orpatashnik/StyleCLIP/blob/adc836d9077b241d92e9ce88da8cb6c01ada2999/models/stylegan2/model.py#L16 But in the latent mapper, I think the input is always (Batch, n_latent, latent_dim)(Maybe I am wrong here), thus the normalization layer seems doesn't do the right normalization. I think for stylespace mapper, this reshape is right: https://github.com/orpatashnik/StyleCLIP/blob/adc836d9077b241d92e9ce88da8cb6c01ada2999/mapper/latent_mappers.py#L98 But for the others, some operations that are not intended may occur as PixelNorm normalize the wrong dim. https://github.com/orpatashnik/StyleCLIP/blob/adc836d9077b241d92e9ce88da8cb6c01ada2999/mapper/latent_mappers.py#L62-L67 Am I right or do I misunderstand something? Look forward to you reply, thanks.