Tensorflow Keras model.loads_weights() breaks on TFElectraModel trained with v4.3.0
## Environment info

- `transformers` version: 4.4.0
- Platform: Linux-4.19.112+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.10
- PyTorch version (GPU?): 1.8.0+cu101 (True)
- Tensorflow version (GPU?): 2.4.1 (True)
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help
@jplu 

## Information

Model I am using TFElectraModel:

The problem arises when using:
* my own modified scripts: see attachment

The tasks I am working on is:
* Information Retrieval on SQuAD dataset

## To reproduce

Steps to reproduce the behavior:

1. train a keras model whith a TFElectraModel as a layer on transformers 4.3.0.
2. switch to transformers 4.4.0
3. call model.load_weights on the keras model

OR:
1. "run all" on attached notebook, wait for weights save after (very short) training
2. change pip tranasformers version to 4.4.0 in the first cell
3. comment model.fit to avoid override of weights
4. restart an run all, wait for model.load_weights to fail

[IR_TFElectraModel.ipynb.zip](https://github.com/huggingface/transformers/files/6151831/IR_TFElectraModel.ipynb.zip)


ValueError: Cannot assign to variable tf_electra_model_1/electra/embeddings/token_type_embeddings/embeddings:0 due to variable shape (2, 128) and value shape (512, 128) are incompatible

## Expected behavior

I would expect the weights to load as it correctly does with both train and load performed on transformers 4.3.0

Anyway keep up with the great job! 🤗 

Thanks for opening this issue!

I have not access to a computer to extract and read you archive, can you, please, share a Colab or copy/paste the code here in this issue. Thanks!<|||||>> Thanks for opening this issue!
> 
> I have not access to a computer to extract and read you archive, can you, please, share a Colab or copy/paste the code here in this issue. Thanks!

yes sure: https://colab.research.google.com/drive/1tHzUMhveYwYkPOCxZwMA80IkVkRCWGRW?usp=sharing

#### To Reproduce

1. "run all" on attached notebook, wait for weights save after (very short) training

2. change pip tranasformers version to 4.4.0 in the first cell
3. comment model.fit to avoid override of weights
4. restart an run all, wait for model.load_weights to fail
<|||||>Sorry, I cannot open your Colab :( There is a restricted access<|||||>Sorry, access at the same link is now allowed. :)<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.