Merge pull request #740 from konstmish:main

PiperOrigin-RevId: 603308692
google-deepmind · Feb 1, 2024 · def75e9 · def75e9
2 parents a1f4ef7 + 82156bb
commit def75e9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/optax/contrib/prodigy.py b/optax/contrib/prodigy.py
@@ -44,7 +44,7 @@ class ProdigyState(NamedTuple):
 
 
 def prodigy(
-    learning_rate: base.ScalarOrSchedule = 0.1,
+    learning_rate: base.ScalarOrSchedule = 1.0,
     betas: tuple[float, float] = (0.9, 0.999),
     beta3: Optional[float] = None,
     eps: float = 1e-8,

diff --git a/optax/contrib/prodigy_test.py b/optax/contrib/prodigy_test.py
@@ -78,7 +78,7 @@ def step(params, state):
     # A no-op change, to verify that tree map works.
     state = _state_utils.tree_map_params(opt, lambda v: v, state)
 
-    for _ in range(10000):
+    for _ in range(100_000):
       params, state = step(params, state)
 
     chex.assert_trees_all_close(params, final_params, rtol=3e-2, atol=3e-2)