Saving/Reloading the Flax T5 model
@patil-suraj @patrickvonplaten 

Hey,
I was able to train the FlaxT5Model using ```run_t5_mlm_flax.py``` on a small custom dataset as per the flax doc. How do I save and reuse the models to get the last hidden states? Would the following code work?

```
auto_tokenizer = AutoTokenizer.from_pretrained("my_output_folder")
model = FlaxT5Model.from_pretrained('my_output_folder')
input_ids = auto_tokenizer("my encoder input", return_tensors="np").input_ids
decoder_input_ids = auto_tokenizer("my decoder input", return_tensors="np").input_ids
outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids)
last_hidden_states = outputs.last_hidden_state
print(numpy.array(last_hidden_states))
```



You can save the model with `model.save_pretrained('my_saved_folder')`