From cd022ff593738c8d9c0b411e294008690b22e1ba Mon Sep 17 00:00:00 2001 From: Zohten Date: Sun, 18 Jul 2021 18:00:21 +0200 Subject: [PATCH] fix stable weight decay for madgrad --- ranger21/ranger21.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ranger21/ranger21.py b/ranger21/ranger21.py index bdf7b82..fed39da 100644 --- a/ranger21/ranger21.py +++ b/ranger21/ranger21.py @@ -702,7 +702,10 @@ class Ranger21(TO.Optimizer): raise ValueError("failed to set param size") # stable weight decay - variance_normalized = math.sqrt(variance_ma_sum / param_size) + if self.use_madgrad: + variance_normalized = torch.pow(variance_ma_sum / param_size, 1/3) + else: + variance_normalized = math.sqrt(variance_ma_sum / param_size) # variance_mean = variance_ma_sum / param_size if math.isnan(variance_normalized): raise RuntimeError("hit nan for variance_normalized")