`model_name_or_path` does not seem to load in previously trained checkpoints
Author: MalteHB
Pull request: no

## 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?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: Using TPU
- Using distributed or parallel set-up in script?: Yes

## Information

Model I am using is RoBERTa, and it is a part of the flax-community week.

I am trying to load a previously trained model checkpoint by setting the 'model_name_or_path' flag into a MLM script, which can be found [here](https://huggingface.co/flax-community/roberta-large-scandi/blob/main/src/run_mlm_flax_stream.py), but seems that the model is initialized with new weights...

## Expected behavior
Seeing that the model training loss would continue from where it stopped, and not seeing that the new model metrics simply mimicked the already trained metrics.

Comments:
Can you post a code snippet or make a Colab to reproduce the error?

@MalteHB it would be nice if you could provide an exact code snippet that we can copy paste to reproduce the error. Otherwise I don't really know what code you've run. I tried re-starting to train from a pretrained checkpoint and it works just fine on my side.

@NielsRogge @patrickvonplaten yes of course, sorry! 

As mentioned we are using a modified version of the `run_mlm_flax_stream.py` script which you can find [here](https://huggingface.co/flax-community/roberta-large-scandi/blob/main/src/run_mlm_flax_stream.py), and the code used to run the script is, where `"/home/Z6HJB/roberta-large-scandi/roberta-base-pretrained-scandinavian/"` is a directory with a `config.json`and a `flax_model.msgpack`: 
```
export MODEL_DIR=/home/Z6HJB/roberta-large-scandi/roberta-base-pretrained-scandinavian/

source /home/Z6HJB/test/bin/activate

python3 ./src/run_mlm_flax_stream.py \
    --model_name_or_path="/home/Z6HJB/roberta-large-scandi/roberta-base-pretrained-scandinavian/" \
    --output_dir="/home/Z6HJB/roberta-large-scandi/model_continued2" \
    --tokenizer_name="${MODEL_DIR}" \
    --dataset_name="mc4" \
    --dataset_config_name="unshuffled_deduplicated_en" \
    --max_seq_length="128" \
    --per_device_train_batch_size="128" \
    --per_device_eval_batch_size="128" \
    --learning_rate="3e-4" \
    --warmup_steps="1000" \
    --overwrite_output_dir \
    --adam_beta1="0.9" \
    --adam_beta2="0.98" \
    --num_train_steps="1000000" \
    --num_eval_samples="5000" \
    --save_steps="1000" \
    --logging_steps="25" \
    --eval_steps="1000" \
    --push_to_hub \
    #--config_name="${MODEL_DIR}" \
    #--model_type="roberta" \
```
Let me know if this suffices or if you need more!

I might be busy for the rest of the day since I have a football match to watch 🇩🇰 🇩🇰 🇩🇰 🇩🇰 🇩🇰 🇩🇰 

I figured it out this line: https://huggingface.co/flax-community/roberta-large-scandi/blob/main/src/run_mlm_flax_stream.py#L439

Forces the model to be reinitialized from scratch everytime. I see that the official script also does that -> I'll open a PR to fix it in the official script and then you shoudl be able to copy paste from it :-) 

What an awesome guy, you are @patrickvonplaten! Thank you so much!