schedule = optax.linear_schedule(
    init_value=0.0,
    end_value=0.001,
    transition_steps=1000
)

        __
        __