rms norm test succeed

microsoft · Aug 15, 2023 · 804ead9 · 804ead9
1 parent 15814d3
commit 804ead9
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/python/unittest/test_layernorm.py b/python/unittest/test_layernorm.py
@@ -46,10 +46,8 @@ def test_layernorm_internal(
     elif layer_norm_type == "rmsnorm":
         gt = (
             torch_input
-            * torch.sqrt(
-                torch.mean(torch_input**2, dim=-1, keepdim=True) + 1e-5
-            ).numpy()
-        )
+            * torch.rsqrt(torch_input.pow(2).mean(-1, keepdim=True) + 1e-5)
+        ).numpy()
 
     # test if the result is correct
     max_abs_error = np.max(np.abs(output_tensor_host - gt))