update: LARS

Author: kozistr

pytorch_optimizer/lars.py

@@ -27,18 +27,21 @@ def __init__(
         weight_decay: float = 0.0,
         momentum: float = 0.9,
         trust_coefficient: float = 0.001,
+        eps: float = 1e-6,
     ):
         """LARS optimizer, no rate scaling or weight decay for parameters <= 1D
         :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups
         :param lr: float. learning rate
         :param weight_decay: float. weight decay (L2 penalty)
         :param momentum: float. momentum
         :param trust_coefficient: float. trust_coefficient
+        :param eps: float. epsilon
         """
         self.lr = lr
         self.weight_decay = weight_decay
         self.momentum = momentum
         self.trust_coefficient = trust_coefficient
+        self.eps = eps
 
         self.check_valid_parameters()
 
@@ -59,6 +62,8 @@ def check_valid_parameters(self):
             raise ValueError(f'Invalid momentum : {self.momentum}')
         if self.trust_coefficient < 0.0:
             raise ValueError(f'Invalid trust_coefficient : {self.trust_coefficient}')
+        if self.eps < 0.0:
+            raise ValueError(f'Invalid eps : {self.eps}')
 
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS:
@@ -84,7 +89,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                     q = torch.where(
                         param_norm > 0.0,
-                        torch.where(update_norm > 0, (g['trust_coefficient'] * param_norm / update_norm), one),
+                        torch.where(update_norm > 0.0, (g['trust_coefficient'] * param_norm / update_norm), one),
                         one,
                     )
                     dp = dp.mul(q)