[Fix] Use a smaller eps value when adding to exp_avg_sq in apply_ams_bound() (#398)

* docs: v3.6.2 changelog

* update: exp_avg_sq_eps

* docs: v3.6.2 changelog

Author: kozistr

docs/changelogs/v3.6.2.md

@@ -4,3 +4,7 @@
 
 * Implement `AdaMuon` optimizer. (#394, #395)
     * [Adaptive Muon Optimizer](https://arxiv.org/abs/2507.11005v1)
+
+### Fix
+
+* Adjust the value of `eps` to the fixed value `1e-15` when adding to `exp_avg_sq`. (#397, #398)

pytorch_optimizer/base/optimizer.py

@@ -154,23 +154,28 @@ def apply_weight_decay(
 
     @staticmethod
     def apply_ams_bound(
-        ams_bound: bool, exp_avg_sq: torch.Tensor, max_exp_avg_sq: Optional[torch.Tensor], eps: float
+        ams_bound: bool,
+        exp_avg_sq: torch.Tensor,
+        max_exp_avg_sq: Optional[torch.Tensor],
+        eps: float,
+        exp_avg_sq_eps: float = 1e-15,
     ) -> torch.Tensor:
         r"""Apply AMSBound variant.
 
         :param ams_bound: bool. whether to apply AMSBound.
         :param exp_avg_sq: torch.Tensor. exp_avg_sq.
         :param max_exp_avg_sq: Optional[torch.Tensor]. max_exp_avg_sq.
         :param eps: float. epsilon.
+        :param exp_avg_sq_eps: float. eps value for numerical stability for exp_avg_sq.
         """
         if ams_bound:
             if torch.is_complex(max_exp_avg_sq):
                 max_exp_avg_sq = torch.view_as_real(max_exp_avg_sq)
 
             torch.maximum(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
-            de_nom = max_exp_avg_sq.add(eps)
+            de_nom = max_exp_avg_sq.add(exp_avg_sq_eps)
         else:
-            de_nom = exp_avg_sq.add(eps)
+            de_nom = exp_avg_sq.add(exp_avg_sq_eps)
 
         return de_nom.sqrt_().add_(eps)