fix image attention layernorm

lightweight_gan/lightweight_gan.py

@@ -156,7 +156,17 @@ def forward(self, x):
         fn = self.fn if random() < self.prob else self.fn_else
         return fn(x)
 
-ChanNorm = partial(nn.InstanceNorm2d, affine = True)
+class ChanNorm(nn.Module):
+    def __init__(self, dim, eps = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
+        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()
+        mean = torch.mean(x, dim = 1, keepdim = True)
+        return (x - mean) / (std + self.eps) * self.g + self.b
 
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):

lightweight_gan/version.py

@@ -1 +1 @@
-__version__ = '0.20.3'
+__version__ = '0.20.4'