update: ranger21

Author: kozistr

pytorch_optimizer/ranger21.py

@@ -241,11 +241,11 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 variance_ma.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                 variance_ma_sum += (variance_ma / bias_correction2).sum()
 
-        if not self.param_size:
+        if self.param_size == 0:
             self.param_size = param_size
 
         # stable weight decay
-        variance_normalized = math.sqrt(variance_ma_sum / param_size)
+        variance_normalized = math.sqrt(variance_ma_sum / self.param_size)
         if math.isnan(variance_normalized):
             raise RuntimeError('hit nan for variance_normalized')