added missing torch addmm with mm

Author: mkhona-nvidia

emerging_optimizers/orthogonalized_optimizers/spectral_clipping_utils.py

@@ -71,7 +71,9 @@ def spectral_hardcap(X: torch.Tensor, beta: float = 1.0) -> torch.Tensor:
     OX = newton_schulz(X, steps=8, coefficient_type="polar_express")
     aX = torch.add(beta * OX, X, alpha=-1)
     result = torch.add(beta * OX, X)
-    result = torch.add(result, aX @ newton_schulz(aX, steps=8, coefficient_type="polar_express").T @ OX, alpha=-1)
+    result = torch.addmm(
+        result, aX, torch.mm(newton_schulz(aX, steps=8, coefficient_type="polar_express").T, OX), alpha=-1
+    )
     result = result * 0.5
     if needs_transpose:
         result = result.T