pow(n,2) => n*n

Summary: pow(n,2) => n*n

Reviewed By: manuelcandales

Differential Revision: D54142779

fbshipit-source-id: 91185e07e9fe040df090756d78be2cf5e87f6e1d

Author: Michael Gschwind

examples/models/llama2/model.py

@@ -62,7 +62,7 @@ def _norm(self, x):
             torch.Tensor: The normalized tensor.
 
         """
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return x * torch.rsqrt((x * x).mean(-1, keepdim=True) + self.eps)
 
     def forward(self, x):
         """