poor GPU utilization on the deep learning examples about jaxopt HOT 4 OPEN

Traceback (most recent call last):
  File "examples/deep_learning/flax_image_classif.py", line 199, in <module>
    app.run(main)
  File "/home/pedregosa/anaconda3/lib/python3.8/site-packages/absl/app.py", line 300, in run
    _run_main(main, args)
  File "/home/pedregosa/anaconda3/lib/python3.8/site-packages/absl/app.py", line 251, in _run_main
    sys.exit(main(argv))
  File "examples/deep_learning/flax_image_classif.py", line 188, in main
    params, state = jax.jit(solver.update)(params=params, state=state,
  File "/home/pedregosa/dev/jaxopt/jaxopt/_src/optax_wrapper.py", line 120, in update
    params, state = self.pre_update(params, state, *args, **kwargs)
  File "examples/deep_learning/flax_image_classif.py", line 145, in print_accuracy
    if state.iter_num % 10 == 0:
jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
The problem arose with the `bool` function. 
While tracing the function update at /home/pedregosa/dev/jaxopt/jaxopt/_src/optax_wrapper.py:104 for jit, this concrete value was not available in Python because it depends on the value of the argument 'state'.

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError