Correct LayerNormalization

newlightlw · Jul 15, 2017 · b40c2da · b40c2da
1 parent ca8c7c0
commit b40c2da
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/transformer/Modules.py b/transformer/Modules.py
@@ -47,8 +47,8 @@ def forward(self, z):
  if z.size(1) == 1:
  return z
 
- mu = torch.mean(z, dim=1)
- sigma = torch.std(z, dim=1)
+ mu = torch.mean(z, dim=-1)
+ sigma = torch.std(z, dim=-1)
  ln_out = (z - mu.expand_as(z)) / (sigma.expand_as(z) + self.eps)
  ln_out = ln_out * self.a_2.expand_as(ln_out) + self.b_2.expand_as(ln_out)