[BUG] QML JAX and Optax demo does not use the correct definition of optimizer

[Gopal-Dahale]

Issue description

In the demo "How to optimize a QML model using JAX and Optax", the update_step_jit does not use the correct definition of the optimizer that is defined in optimization_jit (Line 170) as it is not passed as an argument to the update_step_jit function. It possibly uses the earlier definition of the optimizer (Line 123).

• Expected behavior: In principle, the optimizer should get passed to the update_step_jit function (Line 165-189)

• Actual behavior: The code works fine but uses the earlier definition of the optimizer.

• System information: Using standard google colab.

[KetpuntoG]

Thanks for opening the issue, I look at it!

[KetpuntoG]

Solved 😄

[KetpuntoG]

The changed will be deployed in a few hours