Unexpected behavior when using TFRoberta model inside tf.keras model
## 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.

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.