Invalid argument error with TFRoberta on GLUE
## 🐛 Bug

I am unable to get TFRoberta working on the GLUE benchmark.

Model I am using (Bert, XLNet....): TFRoberta

Language I am using the model on (English, Chinese....): English

The problem arise when using:
* [ ] the official example scripts: (give details)
* [x] my own modified scripts: (give details)

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: *MRPC*
* [ ] my own task or dataset: (give details)

## To Reproduce

Steps to reproduce the behavior:

1. Copy the following script into a file name `run_glue_task_tf.py`:
```python
import tensorflow as tf
import tensorflow.keras as keras

import tensorflow_datasets.public_api as tfds

import transformers

from absl import app
from absl import flags
from absl import logging

FLAGS = flags.FLAGS
flags.DEFINE_multi_string("tasks", None, "One or more tasks to be used for pretraining")

flags.DEFINE_integer('num_epochs', 3, 'Number of epochs to train')
flags.DEFINE_integer('batch_size', 32, 'Batch size to use for training')
flags.DEFINE_integer('eval_batch_size', 64, 'Batch size to use when evaluating validation/test sets')
flags.DEFINE_boolean('use_xla', False, 'Enable XLA optimization')
flags.DEFINE_boolean('use_amp', False, 'Enable AMP optimization')
flags.DEFINE_integer('max_seq_len', 128, 'Maximum sequence length')
flags.DEFINE_string('model_name', 'bert-base-cased', 'Name of pretrained transformer model to load')


def configure_tf():
    logging.info(('Enabling' if FLAGS.use_xla else 'Disabling') + ' XLA optimization')
    tf.config.optimizer.set_jit(FLAGS.use_xla)
    logging.info(('Enabling' if FLAGS.use_xla else 'Disabling') + ' auto mixed precision (AMP)')
    tf.config.optimizer.set_experimental_options({'auto_mixed_precision': FLAGS.use_amp})


def get_model():
    logging.info('Loading pre-trained TF model from %s', FLAGS.model_name)
    model = transformers.TFAutoModelForSequenceClassification.from_pretrained(FLAGS.model_name)

    opt = keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08)
    if FLAGS.use_amp:
        logging.debug('Enabling loss scaling')
        opt = keras.mixed_precision.experimental.LossScaleOptimizer(opt, 'dynamic')

    loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)

    metric = keras.metrics.SparseCategoricalAccuracy('accuracy')

    model.compile(optimizer=opt, loss=loss, metrics=[metric])

    return model


def run_task(task: str, model: keras.Model, tokenizer):
    data, info = tfds.load(task, with_info=True)

    if task.startswith('glue'):
        glue_task = task[len('glue/'):]

        def load_features(split: tfds.core.splits.NamedSplit):
            logging.debug('Converting %s.%s to features', task, split)
            is_xlnet: bool = 'xlnet' in model.name.lower()
            return transformers.glue_convert_examples_to_features(examples=data[split],
                                                                  tokenizer=tokenizer,
                                                                  max_length=FLAGS.max_seq_len,
                                                                  output_mode=transformers.glue_output_modes[glue_task],
                                                                  pad_on_left=is_xlnet,  # Pad on the left for XLNet
                                                                  pad_token=tokenizer.convert_tokens_to_ids(
                                                                      [tokenizer.pad_token])[0],
                                                                  pad_token_segment_id=4 if is_xlnet else 0,
                                                                  task=glue_task)

        train = load_features(tfds.Split.TRAIN)
        valid = load_features(tfds.Split.VALIDATION)
    else:
        raise ValueError('Unsupported task: %s' % task)

    train = train.shuffle(128).batch(FLAGS.batch_size).repeat(FLAGS.num_epochs)
    valid = valid.batch(FLAGS.eval_batch_size).take(FLAGS.eval_batch_size)

    logging.info('Training %s on %s...', model.name, task)
    history = model.fit(x=train, validation_data=valid)

    print(task, 'Performance:')
    print(history)


def main(argv):
    del argv  # Unused.

    logging.debug('Loading tokenizer from %s', FLAGS.model_name)
    tokenizer = transformers.AutoTokenizer.from_pretrained(FLAGS.model_name)
    model = get_model()

    for task in FLAGS.tasks:
        print('-' * 20, task, '-' * 20)
        run_task(task, model, tokenizer)


if __name__ == '__main__':
    app.run(main)
```
2. Run the following command:
```bash
$ python run_glue_task_tf.py --tasks glue/mrpc --model_name roberta-base
```
3. Experience the following error:
```
I1101 14:04:14.134916 46912496418432 run_pretraining.py:76] Training tf_roberta_for_sequence_classification on glue/mrpc...
/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/framework/indexed_slices.py:424: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.
  "Converting sparse IndexedSlices to a dense Tensor of unknown shape. "
WARNING:tensorflow:Gradients do not exist for variables ['tf_roberta_for_sequence_classification/roberta/pooler/dense/kernel:0', 'tf_roberta_for_sequence_classification/roberta/pooler/dense/bias:0'] when minimizing the loss.
W1101 14:04:25.996359 46912496418432 optimizer_v2.py:1029] Gradients do not exist for variables ['tf_roberta_for_sequence_classification/roberta/pooler/dense/kernel:0', 'tf_roberta_for_sequence_classification/roberta/pooler/dense/bias:0'] when minimizing the loss.
/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/framework/indexed_slices.py:424: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.
  "Converting sparse IndexedSlices to a dense Tensor of unknown shape. "
WARNING:tensorflow:Gradients do not exist for variables ['tf_roberta_for_sequence_classification/roberta/pooler/dense/kernel:0', 'tf_roberta_for_sequence_classification/roberta/pooler/dense/bias:0'] when minimizing the loss.
W1101 14:04:36.157998 46912496418432 optimizer_v2.py:1029] Gradients do not exist for variables ['tf_roberta_for_sequence_classification/roberta/pooler/dense/kernel:0', 'tf_roberta_for_sequence_classification/roberta/pooler/dense/bias:0'] when minimizing the loss.
2019-11-01 14:04:45.689286: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Invalid argument: indices[28,20] = 1 is not in [0, 1)
         [[{{node tf_roberta_for_sequence_classification/roberta/embeddings/token_type_embeddings/embedding_lookup}}]]
2019-11-01 14:04:45.900194: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Invalid argument: indices[28,20] = 1 is not in [0, 1)
         [[{{node tf_roberta_for_sequence_classification/roberta/embeddings/token_type_embeddings/embedding_lookup}}]]
         [[Shape/_8]]
      1/Unknown - 28s 28s/stepTraceback (most recent call last):
  File "run_pretraining.py", line 96, in <module>
    app.run(main)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/absl/app.py", line 299, in run
    _run_main(main, args)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/absl/app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "run_pretraining.py", line 92, in main
    run_task(task, model, tokenizer)
  File "run_pretraining.py", line 77, in run_task
    history = model.fit(x=train, validation_data=valid)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py", line 728, in fit
    use_multiprocessing=use_multiprocessing)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 324, in fit
    total_epochs=epochs)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 123, in run_one_epoch
    batch_outs = execution_function(iterator)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2_utils.py", line 86, in execution_function
    distributed_function(input_fn))
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py", line 457, in __call__
    result = self._call(*args, **kwds)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py", line 520, in _call
    return self._stateless_fn(*args, **kwds)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1823, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1141, in _filtered_call
    self.captured_inputs)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1224, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 511, in call
    ctx=ctx)
  File "/data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/eager/execute.py", line 67, in quick_execute
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError:  indices[28,20] = 1 is not in [0, 1)
         [[node tf_roberta_for_sequence_classification/roberta/embeddings/token_type_embeddings/embedding_lookup (defined at /data/conda/envs/transformers/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1751) ]] [Op:__inference_distributed_function_58202]

Function call stack:
distributed_function

2019-11-01 14:04:47.519198: W tensorflow/core/kernels/data/generator_dataset_op.cc:102] Error occurred when finalizing GeneratorDataset iterator: Failed precondition: Python interpreter state is not initialized. The process may be terminated.
         [[{{node PyFunc}}]]
```

## Expected behavior

The script should train on GLUE when using TFRoberta just as it trains when using TFBert.

## Environment

OS: CentOS-7
Python version: 3.6.9
PyTorch version: 1.2.0
PyTorch Transformers version (or branch): 2.1.1 (master from git)
Using GPU ? Yes
Distributed of parallel setup ? No
Any other relevant information:

## Additional context

I have encountered this embedding error and I found that wrapping the model in a tf.distribute strategy fixes it (extremely odd...)

You can wrap the model in `tf.distribute.OneDeviceStrategy` and it will work.<|||||>Did you ever figure it out? <|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
