How to get params of optimizer?

[arihantbansal]

Is there any way to get parameters of the optimizer/network at a step given the optimizer state? Similar to get_params(opt_state) which exists in jax.example_libraries.optimizers

[pharringtonp19]

@arihantbansal

I think you are looking for the following two lines of code.

updates, opt_state = optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)

At each iteration

1. Update the optimizer state optimizer.update(grads, opt_state)
2. Update the parameters of the neural network optax.apply_updates(params, updates)