Not able to load the custom model after training in Hugging Face
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version:
- Platform: colab

### Who can help

 @patil-suraj @patrickvonplaten 

Models:
Used FlaxGPT2Module as base class and built FlaxGPT2ForMutlipleChoice.After training is done when loading the saved weights in hugging face hub there seems to be an error.
Model hub:

-Hugging Face Hub:https://huggingface.co/Vivek/gpt2-common-sense-reasoning

Error messgae:  unpack(b) received extra data.

## Information

Model I am using (Bert, XLNet ...):FlaxGPT2ForMultipleChoice(custom model)

The problem arises when using: 
1.when i try to load the weights of   saved weights in hugging face hub there seems to be an error.

The tasks I am working on is:
* Dataset: COSMOS QA

Expected Behaviour :To able to load the weights and the configuration without any error

Colab Notebook:https://colab.research.google.com/drive/1C-M1GLMk7jiomXIpngbLZvCNvEJMfJ1K?usp=sharing

## To reproduce
``` python
new_model=FlaxGPT2ForMultipleChoice.from_pretrained('/content/gpt2-common-sense-reasoning',
                 input_shape=(1,4,1), config=config)```

