error in stylegan2_layers.py

[AltoriaDD]

Hello,Thank you for your work.I hope others can notice that when you want dim>2, please make corresponding modifications
style = self.modulation(style) if self.demodulate: style = style * torch.rsqrt(style.pow(2).mean([1], keepdim=True) + 1e-8)