Unexpected behavior when using TFRoberta model inside tf.keras model
Author: amir-ghasemi
Pull request: no

## Environment

  - `transformers` version: tried with both 4.0.0 and 3.5.0
- Platform: Google Colab
- Python version: 3.6.9
- PyTorch version (GPU?):
- Tensorflow version (GPU?): 2.3.0 GPU
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help
@LysandreJik 
@jplu 

## Information

Model I am using (Bert, XLNet ...): TFRoberta

The problem arises when using:
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
Sentence classification 

## To reproduce

I  am trying to import a pretrained TFRoberta model and extend it with a few layers for classification using tensorflow keras. When I directly use transformers model (Method 1), the model trains well and reaches a validation accuracy of 0.93 after 1 epoch. However, when trying to use the model as a layer within a tf.keras model (Method 2), the model can't get above 0.32 accuracy. As far as I can tell based on the documentation, the two approaches should be equivalent. My goal is to get Method 2 working so that I can add more layers to it instead of directly using the logits produced by the transformers' classifier head but I'm stuck at this stage.

```
import tensorflow as tf
from transformers import TFRobertaForSequenceClassification
```

Method 1:

```
model = TFRobertaForSequenceClassification.from_pretrained("roberta-base", num_labels=6)
```

Method 2:

```
input_ids = tf.keras.Input(shape=(128,), dtype='int32')
attention_mask = tf.keras.Input(shape=(128, ), dtype='int32')

transformer = TFRobertaForSequenceClassification.from_pretrained("roberta-base", num_labels=6)

encoded = transformer([input_ids, attention_mask])

logits = encoded[0]

model = tf.keras.models.Model(inputs = [input_ids, attention_mask], outputs = logits) 
```

Rest of the code for either method is identical,

```
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0),
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), 
              metrics=[tf.keras.metrics.SparseCategoricalAccuracy('accuracy')])
```

## Expected behavior
Similar validation loss and accuracy for both methods.


Comments:
Hello !

Thanks for reporting this. Can you provide a colab in order for us to reproduce your use case?

Thanks @jplu . Please see the [Colab notebook](https://colab.research.google.com/drive/1qDpqEc4qbeuoQjVEpDx88dXIXnJ6bwXh?usp=sharing).

As you can see, both model1 and model2 have the exact same number of parameters and are initialized using the same pretrained roberta-base model. Yet, first one trains well and reaches val_accuracy of 0.9350 after one epoch while the second one (using transformers model within a tf.keras model) is stuck.

This issue does not seem to be isolated to TFRoberta. Just tried with TFDistilBertForSequenceClassification and the outcome is similar. Using the transformers model directly works fine whereas embedding it within a tf.keras model (while adding just an input layer and passing the logits directly to output) fails.

@amir-ghasemi Can you try on the master version with this update, and let me know if you still get the issue:
```
input_ids = tf.keras.Input(shape=(128,), dtype='int32')
attention_mask = tf.keras.Input(shape=(128, ), dtype='int32')

transformer = TFRobertaForSequenceClassification.from_pretrained("roberta-base", num_labels=6)
encoded = transformer({"input_ids": input_ids, "attention_mask": attention_mask})
logits = encoded[0]

model = tf.keras.models.Model(inputs = {"input_ids": input_ids, "attention_mask": attention_mask}, outputs = logits)
```

Thanks @jplu ! Tried with the master branch and feeding the input using the dict. That did the trick! Closing the issue.