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:

StyleCLIP/mapper/latent_mappers.py

Line 16 in adc836d

layers = [PixelNorm()]

In the StyleGAN2 mapper, as the input is always (Batch, latent_dim), the normalization layer is right here.

StyleCLIP/models/stylegan2/model.py

Line 16 in adc836d

return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)

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:

StyleCLIP/mapper/latent_mappers.py

Line 98 in adc836d

x_c_res = curr_mapper(x_c.view(x_c.shape[0], -1)).view(x_c.shape)

But for the others, some operations that are not intended may occur as PixelNorm normalize the wrong dim.

StyleCLIP/mapper/latent_mappers.py

Lines 62 to 67 in adc836d

	x_coarse = x[:, :4, :]
	x_medium = x[:, 4:8, :]
	x_fine = x[:, 8:, :]
	if not self.opts.no_coarse_mapper:
	x_coarse = self.course_mapping(x_coarse)

Am I right or do I misunderstand something? Look forward to you reply, thanks.