# Optax
schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,
    peak_value=0.001,
    warmup_steps=1000,
    decay_steps=10000
)

optimizer = optax.adam(learning_rate=schedule)

        __
        __