Fix typing issue in prodigy.py for python 3.9

PiperOrigin-RevId: 589961288
google-deepmind · Dec 11, 2023 · 83360c5 · 83360c5
1 parent 1a7956d
commit 83360c5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optax/contrib/prodigy.py b/optax/contrib/prodigy.py
@@ -46,7 +46,7 @@ class ProdigyState(NamedTuple):
 def prodigy(
     learning_rate: base.ScalarOrSchedule = 0.1,
     betas: tuple[float, float] = (0.9, 0.999),
-    beta3: float | None = None,
+    beta3: Optional[float] = None,
     eps: float = 1e-8,
     estim_lr0: float = 1e-6,
     estim_lr_coef: float = 1.0,