Fix init bug

SymbioticLab · Dec 17, 2023 · d889ac8 · d889ac8
1 parent 442c85c
commit d889ac8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fedscale/utils/optimizer/yogi.py b/fedscale/utils/optimizer/yogi.py
@@ -15,7 +15,7 @@ def __init__(self, eta=1e-2, tau=1e-3, beta=0.9, beta2=0.99):
     def update(self, gradients):
         update_gradients = []
         if not self.v_t:
-            self.v_t = [g**2 for g in gradients]
+            self.v_t = [torch.full_like(g, self.tau) for g in gradients]
             self.m_t = [torch.full_like(g, 0.0) for g in gradients]
 
         for idx, gradient in enumerate(gradients):