[BERT/PyT] fix onnx export (NVIDIA#689)

changlan · Apr 5, 2021 · 1f65809 · 1f65809
1 parent c29efa3
commit 1f65809
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/PyTorch/LanguageModeling/BERT/modeling.py b/PyTorch/LanguageModeling/BERT/modeling.py
@@ -287,7 +287,9 @@ def __init__(self, hidden_size, eps=1e-12):
 
     def forward(self, x):
         u = x.mean(-1, keepdim=True)
-        s = (x - u).pow(2).mean(-1, keepdim=True)
+        s = (x - u)
+        s = s * s
+        s = s.mean(-1, keepdim=True)
         x = (x - u) / torch.sqrt(s + self.variance_epsilon)
         return self.weight * x + self.bias
 
@@ -323,7 +325,9 @@ def forward(self, x):
             x = self.fused_layer_norm(x)
         else:
             u = x.mean(-1, keepdim=True)
-            s = (x - u).pow(2).mean(-1, keepdim=True)
+            s = (x - u)
+            s = s * s
+            s = s.mean(-1, keepdim=True)
             x = (x - u) / torch.sqrt(s + self.eps)
             x = self.weight * x + self.bias
         return x