From 2badbe80d11a4589650daa4a50461ca247bfdd4a Mon Sep 17 00:00:00 2001 From: OptaxDev Date: Tue, 29 Oct 2024 15:59:21 -0700 Subject: [PATCH] Add warmup_constant_schedule to optax.__init__.py. PiperOrigin-RevId: 691192970 --- optax/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/optax/__init__.py b/optax/__init__.py index 688ca13f..7316e017 100644 --- a/optax/__init__.py +++ b/optax/__init__.py @@ -207,6 +207,7 @@ piecewise_interpolate_schedule = schedules.piecewise_interpolate_schedule polynomial_schedule = schedules.polynomial_schedule sgdr_schedule = schedules.sgdr_schedule +warmup_constant_schedule = schedules.warmup_constant_schedule warmup_cosine_decay_schedule = schedules.warmup_cosine_decay_schedule warmup_exponential_decay_schedule = schedules.warmup_exponential_decay_schedule inject_stateful_hyperparams = schedules.inject_stateful_hyperparams