Update tutorial_How_to_optimize_QML_model_using_JAX_and_Optax.py

[KetpuntoG]

Solving issue #1045

As indicated in the issue, the optimizer being used within the update_step_jit function is not the one defined in the optimization_jit function. I tried to pass opt but it gave problems so I simply defined the optimizer outside.

In the case where jit is not being applied, I was able to pass the optimizer, so I updated those functions.
Tagging @josh146 in case there is a better solution in this case

[github-actions]

👋 Hey, looks like you've updated some demos!

🐘 Don't forget to update the dateOfLastModification in the associated metadata files so your changes show up on pennylane.ai.

Please hide this comment once the field(s) are updated. Thanks!

[github-actions]

Thank you for opening this pull request.

You can find the built site at this link.

Deployment Info:

• Pull Request ID: 1052
• Deployment SHA: ce3f10b1c00cd798ddeb77c5f79833e03de11fb8
  (The Deployment SHA refers to the latest commit hash the docs were built from)

Note: It may take several minutes for updates to this pull request to be reflected on the deployed site.

[josh146]

thanks @KetpuntoG

[isaacdevlugt]

LGTM!