fix(chainable): correct nesterov momentum

ClashLuke · Dec 8, 2024 · a7791f0 · a7791f0
1 parent dfc3299
commit a7791f0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/heavyball/chainable.py b/heavyball/chainable.py
@@ -277,13 +277,13 @@ def orthogonalize_update(group, update, grad, param):
 @zero_guard("momentum")
 @no_state
 def nesterov_momentum(group, updates, grads, params, momentum):
-    utils.nesterov_momentum(momentum, updates, utils.get_beta1(group))
+    return utils.nesterov_momentum(momentum, updates, utils.get_beta1(group))
 
 
 @zero_guard("momentum")
 @no_state
 def heavyball_momentum(group, updates, grads, params, momentum):
-    utils.heavyball_momentum(momentum, updates, utils.get_beta1(group))
+    return utils.heavyball_momentum(momentum, updates, utils.get_beta1(group))
 
 
 @zero_guard("exp_avg", "exp_avg_sq")