SVI optimizer scheduling

You can use optax for scheduling (see SVI.optim docs for an example) - this is recommended. Or you can specify schedule like optim.Adam(schedule).