Change initial eigen values in KLShampoo to 1 (#310)

* change init value

Signed-off-by: Hao Wu <skyw@nvidia.com>

* grammar fix

Signed-off-by: Hao Wu <skyw@nvidia.com>

---------

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

emerging_optimizers/shampoo/shampoo.py

@@ -195,8 +195,10 @@ def init_state(
             {
                 "Q_L": torch.eye(m, device=device),
                 "Q_R": torch.eye(n, device=device),
-                "eigvals_L": torch.zeros(m, device=device),
-                "eigvals_R": torch.zeros(n, device=device),
+                # Initialize eigenvalues to one to skip KL correction for the first step.
+                # They will be overwritten immediately afterward.
+                "eigvals_L": torch.ones(m, device=device),
+                "eigvals_R": torch.ones(n, device=device),
             }
         )
         return state