Unpin jax, closes #214

[jobrachem]

The issue in our code originated in the liesel.goose.optim module, specifically in the callback that updates the progress bar.

In the callback, the current loss values were accessed like this:

i = val["while_i"]
loss_train = val["history"]["loss_train"][i]
loss_validation = val["history"]["loss_validation"][i]

Some experimentation revealed that this pattern was the problem. Introducing dedicated "current_loss_train" and "current_loss_validation" carry values solved the issue:

loss_train = val["current_loss_train"]
loss_validation = val["current_loss_validation"]

While I was at it, I fixed two other issues with the progress bar:

1. There was a bug that prevented the argument progress_bar=False from taking effect, i.e. the progress bar couldn't be turned off. This is fixed now.
2. The progress bar displayed max_iter-1 as the total number of iterations. This is fixed now.