Error while performing eval on clm using gpt2 in flax
## Environment info
- `transformers` version: 4.9.0.dev0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): not installed (NA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (tpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

@patrickvonplaten @patil-suraj 



## Information

Model I am using (Bert, XLNet ...):GPT2

The problem arises when using:
* [1 ] the official example scripts: (give details below)
using the examples/flax/language-modeling/run_clm_flax.py
* [ 2] my own modified scripts: (give details below):
a run.sh file of the format -
run.sh --param1 --param2


The tasks I am working on is:
* [ 1] my own task or dataset: (give details below):
* txt file containing rap lyrics starting with <BOS> and ending with <EOS>

## To reproduce

Steps to reproduce the behavior:

1.1.Make a new directory test and change to this directory
2.Add tokenizer.json and config.json from the gpt2 repo from (https://huggingface.co/gpt2/tree/main) to this repository
3.Make a run.sh file of the type run.sh --param1 --param2 and add evaluation parameters such --do_eval and --eval_steps
4.run the file ./run.sh


## Expected behavior
When evaluation occurs you will get the following error:
File "run_clm_flax.py", line 640, in <module>
    main()
  File "run_clm_flax.py", line 609, in main
    eval_metrics = get_metrics(eval_metrics)
  File "/home/anantshankhdhar/RapAiAnant/lib/python3.8/site-packages/flax/training/common_utils.py", line 53, in get_metrics
    return stack_forest(metrics_np)
  File "/home/anantshankhdhar/RapAiAnant/lib/python3.8/site-packages/flax/training/common_utils.py", line 45, in stack_forest
    return jax.tree_multimap(stack_args, *forest)
TypeError: tree_map() missing 1 required positional argument: 'tree'

Hey @AnantShankhdhar - could you please provide the `run.sh` file?<|||||>The error was because the eval batch size was very high 