Problem while using checkpoints.restore_checkpoint with gradio

Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`:
- Python version:
- GPU/TPU model and memory:
- CUDA version (if applicable):

### Problem you have encountered:

Throws exception when using **flax.training.checkpoints.restore_checkpoint** within a module wrapped with using gradio. After some sanity checks and from the logs, I feel that this happens because of the use of parallelization within .restore_checkpoint that is not supported by gradio.

### Logs, error messages, etc:

### Steps to reproduce:

```
T5_PRETRAINED_PATH="path_to_the_weight"
PARAM_DTYPE=jnp.float32
PARAMS_LANG = checkpoints.restore_checkpoint(T5_PRETRAINED_PATH, target=None,parallel=False)
PARAMS_LANG=jax.tree_map(lambda x:x.astype(PARAM_DTYPE),PARAMS_LANG)

with gr.Blocks() as demo:

with gr.Column():
input_file_obj = gr.File(label="File upload")
Num_of_samples=gr.Text(label="Number of samples",value=100)
submit_button = gr.Button("Generate")
input_load=[input_file_obj,Num_of_samples]
with gr.Column():
file_out_obj = gr.File(label="Output File")
submit_button.click(pass_through_func,
inputs=input_load,
outputs=file_out_obj)
demo.launch()

```