Numpy arrays used instead of jax array in example
### System Info

The example[ here](https://github.com/huggingface/transformers/blob/main/examples/flax/text-classification/run_flax_glue.py#L286) should yield jax device arrays instead of numpy arrays. 

@patil-suraj

### Who can help?

_No response_

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

https://github.com/huggingface/transformers/blob/main/examples/flax/text-classification/run_flax_glue.py#L286

### Expected behavior

Should return batches with jax device arrays
Answered in https://github.com/huggingface/transformers/issues/18289#issuecomment-1198216304