more stable chan layernorm

Author: lucidrains

stylegan2_pytorch/stylegan2_pytorch.py

@@ -95,9 +95,9 @@ def __init__(self, dim, eps = 1e-5):
         self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
 
     def forward(self, x):
-        std = torch.var(x, dim = 1, unbiased = False, keepdim = True).sqrt()
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
         mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) / (std + self.eps) * self.g + self.b
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):

stylegan2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.8.7'
+__version__ = '1.8.8'