fix wd for sparse

DenisVorotyntsev · Oct 1, 2022 · d95c8dd · d95c8dd
1 parent 02e6624
commit d95c8dd
Show file tree

Hide file tree

Showing 3 changed files with 168 additions and 181 deletions.
diff --git a/notebooks/test_adan.ipynb b/notebooks/test_adan.ipynb
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
     name="adan-tensorflow",
     packages=find_packages(exclude=[]),
-    version="1.0.0",
+    version="1.0.1",
     license="MIT",
     description="Adan - (ADAptive Nesterov momentum algorithm) Optimizer in Tensorflow",
     author="Denis Vorotyntsev",

diff --git a/tf_adan/adan.py b/tf_adan/adan.py
@@ -211,10 +211,11 @@ def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
             (exp_avg / bias_correction1 + beta2 * exp_avg_diff / bias_correction2)
         ) / denom
 
-        var_update = var.assign_sub(var_update * lr) / (1 + lr * weight_decay)
+        var_updated = var - var_update * lr
+        var_updated = var_updated / (1 + lr * weight_decay)
         return tf.group(
             *[
-                var_update,
+                var.assign(var_updated, use_locking=self._use_locking),
                 exp_avg,
                 exp_avg_diff,
                 exp_avg_sq,