Jax/Flax Text-Classification Examples are not working...
## Environment info
- `transformers` version: `pip install -q git+https://github.com/huggingface/transformers.git`
- Platform: Colab | Kaggle Code 
- Python version: Python3.6 | Python3.7
- Jax version (TPU): 0.2.17
- Using distributed or parallel set-up in script?: Yes (8 TPUv3 given by Colab or Kaggle)

### Who can help
- jax/flax examples: @patil-suraj @patrickvonplaten 

## Information

Model I am using RoBerta, Bert through [jax/flax examples](https://github.com/huggingface/transformers/tree/master/examples/flax#jaxflax-examples) for regression tasks.

The problem arises when using:
* [ ] the python script: [run_flax_glue.py](https://github.com/huggingface/transformers/blob/master/examples/flax/text-classification/run_flax_glue.py)
 => EDIT: `Not working inside colab or kaggle environments`. Please **REPRODUCE** using [colab notebook](https://colab.research.google.com/drive/11-Zj8MRxdgsWIqoX_oixbqk5MdCf4MpW?usp=sharing)
* [ ] the python script: [run_flax_glue.py](https://github.com/huggingface/transformers/blob/master/examples/flax/text-classification/run_flax_glue.py) has a simple bug when used for **custom regression task**
=> EDIT please **REPRODUCE** using [colab notebook](https://colab.research.google.com/drive/1-RCcDgiQ6EGswHaLUOfX_l_0le6naG9B?usp=sharing)
* [ ] the colab notebook: [text_classification_flax.ipynb](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification_flax.ipynb) => EDIT: This is working fine!
The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: Regression task - stsb 
EDIT: and Custom regression task.
@patrickvonplaten validated this is working in TPUv3-8 but it is not working in Colab & Kaggle env

## Expected behavior
For python script, it gets stuck with this output:
```
07/19/2021 18:50:23 - INFO - __main__ - ===== Starting training (20 epochs) =====
/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:387: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.
  "jax.host_count has been renamed to jax.process_count. This alias "
/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:374: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update your code.
  "jax.host_id has been renamed to jax.process_index. This alias "
07/19/2021 18:50:23 - INFO - __main__ - Epoch 1
07/19/2021 18:50:23 - INFO - __main__ -   Training...
2021-07-19 18:53:12.601711: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:55] 
********************************
Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
Compiling module pmap_train_step.64578
********************************
```

For notebook, it gets stuck here:
![Screen Shot 2021-07-19 at 10 48 58 PM](https://user-images.githubusercontent.com/43239645/126260614-987938d1-236c-4f37-89ca-a10ab2958682.png)

There was also a little bug in [python script](https://github.com/huggingface/transformers/blob/master/examples/flax/text-classification/run_flax_glue.py#L358) where label_list is not referenced.
We can declare as None in [L297](https://github.com/huggingface/transformers/blob/master/examples/flax/text-classification/run_flax_glue.py#L297)
Then check if it is assigned in [L357](https://github.com/huggingface/transformers/blob/master/examples/flax/text-classification/run_flax_glue.py#L357)

EDIT:
This error occurred when I don't try with glue task (stsb) but with custom data.<|||||><del>Do you want to open a PR to fix it maybe? :-) </del>

Actually I cannot reproduce the error, when running:

```
python run_flax_glue.py \
  --model_name_or_path roberta-base \
  --task_name stsb \
  --max_length 128 \
  --learning_rate 2e-5 \
  --num_train_epochs 1 \
  --per_device_train_batch_size 4 \
  --output_dir "./output" \
```
 on TPUv3-8 the training finished in ~2 minutes<|||||>@patrickvonplaten 
Sorry for the confusion.
For the code getting stuck, Please take a look at this [one](https://colab.research.google.com/drive/11-Zj8MRxdgsWIqoX_oixbqk5MdCf4MpW?usp=sharing)
For the not referenced error, I also confused and it was with the custom data. Please take a look at this [one](https://colab.research.google.com/drive/1-RCcDgiQ6EGswHaLUOfX_l_0le6naG9B?usp=sharing)<|||||>@patrickvonplaten 
colab example [text_classification_flax.ipynb](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification_flax.ipynb) is working fine. 
So I am trying to find out why it gets stuck when I use [python script inside colab](https://colab.research.google.com/drive/11-Zj8MRxdgsWIqoX_oixbqk5MdCf4MpW?usp=sharing).<|||||>Hey @bayartsogt-ya, 

Note that the Python script is not really made for colab and we recommend using the [text_classification_flax.ipynb](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification_flax.ipynb) instead.
Could you try using the official text classification notebook instead? <|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.