Unable to train sequence classification task using TFTrainer
## 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: 4.1.1
- Platform: Linux-4.15.0-123-generic-x86_64-with-debian-buster-sid
- Python version: 3.7.9
- PyTorch version (GPU?): not installed (NA)
- Tensorflow version (GPU?): 2.4.0 (False)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
 
 
 -->
Trainer: @sgugger 
Tensorflow:  @jplu

## Information

Model I am using (Bert, XLNet ...): distilbert-base-cased

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

The following is how I load the model and the trainer
```{python}
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from datasets import load_from_disk
import tensorflow as tf
from transformers import TFAutoModelForSequenceClassification, AutoTokenizer
from transformers import TFTrainingArguments, TFTrainer

# Load dataset
def load_dataset(data_dir, split="train", batch_size=32, shuffle=100):
    dataset = load_from_disk(data_dir)[split]

    label_type = tf.int32
    input_names = ["input_ids", "attention_mask", "token_type_ids"]

    def gen():
        for ex in dataset:
            d = {k: v for k, v in ex.items() if v is not None}
            label = d.pop("tag")
            yield (d, label)

    tf_dataset = tf.data.Dataset.from_generator(
        gen,
        ({k: tf.int32 for k in input_names}, label_type),
        ({k: tf.TensorShape([None]) for k in input_names}, tf.TensorShape([])),
    )
    tf_dataset = tf_dataset.apply(tf.data.experimental.assert_cardinality(len(dataset)))
    return tf_dataset


# Load the model
def load_model(name="distilbert-base-cased", num_labels=11, learning_rate=3e-5):
    tokenizer = AutoTokenizer.from_pretrained(name)
    tokenizer.add_special_tokens({"bos_token": "<s>", "eos_token": "</s>"})
    model = TFAutoModelForSequenceClassification.from_pretrained(
        name, num_labels=num_labels
    )
    model.resize_token_embeddings(len(tokenizer))
    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    model.compile(optimizer=optimizer, loss=loss)
    model.summary()
    return model


#metrics
def compute_metrics(pred):
    labels = pred.label_ids
    predictions = pred.predictions.argmax(-1)
    precision, recall, f1, _ = precision_recall_fscore_support(
        labels, predictions, average="weighted"
    )
    acc = accuracy_score(labels, predictions)
    return {"accuracy": acc, "f1": f1, "precision": precision, "recall": recall}


# train model
def train_model(
    model_args, data_dir, model_dir, logs_dir, batch_size=32, num_epochs=10
):
    training_args = TFTrainingArguments(
        output_dir=model_dir,
        num_train_epochs=num_epochs,
        do_train=True,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size * 2,
        evaluation_strategy="steps",
        warmup_steps=500,
        weight_decay=0.01,
        logging_dir=logs_dir,
        dataloader_num_workers=15,
    )
    datasets = {
        "train": load_dataset(data_dir=data_dir, split="train", batch_size=batch_size),
        "val": load_dataset(
            data_dir=data_dir, split="validation", batch_size=batch_size
        ),
    }
    
    model = load_model(**model_args)

    trainer = TFTrainer(
        model=model,
        args=training_args,
        train_dataset=datasets["train"],
        eval_dataset=datasets["val"],
        compute_metrics=compute_metrics,
    )
    trainer.train()
    return trainer
```
The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)
A multi-class classification task to classify a sentence into one of 11 known categories

## To reproduce

Steps to reproduce the behavior:

1. A classification task with more the 2 categories - i.e. num_labels > 2.
2. Use pretrained distill bert model for sequence classification 
3. Load the dataset and finetune the model with TFTrainer.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

Here is the stack trace from the error.
```
ValueError: in user code:

    /root/.local/share/virtualenvs/ai-dialogue-acts-classifier-ZQsBEC0q/lib/python3.7/site-packages/transformers/trainer_tf.py:678 distributed_training_steps  *
        self.args.strategy.run(self.apply_gradients, inputs)
    /root/.local/share/virtualenvs/ai-dialogue-acts-classifier-ZQsBEC0q/lib/python3.7/site-packages/transformers/trainer_tf.py:641 apply_gradients  *
        self.optimizer.apply_gradients(list(zip(gradients, self.model.trainable_variables)))
    /root/.local/share/virtualenvs/ai-dialogue-acts-classifier-ZQsBEC0q/lib/python3.7/site-packages/transformers/optimization_tf.py:232 apply_gradients  *
        return super(AdamWeightDecay, self).apply_gradients(zip(grads, tvars), name=name, **kwargs)
    /root/.local/share/virtualenvs/ai-dialogue-acts-classifier-ZQsBEC0q/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:604 apply_gradients  **
        self._create_all_weights(var_list)
    /root/.local/share/virtualenvs/ai-dialogue-acts-classifier-ZQsBEC0q/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:783 _create_all_weights
        self._create_slots(var_list)
    /root/.local/share/virtualenvs/ai-dialogue-acts-classifier-ZQsBEC0q/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/adam.py:127 _create_slots
        self.add_slot(var, 'm')
    /root/.local/share/virtualenvs/ai-dialogue-acts-classifier-ZQsBEC0q/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:844 add_slot
        .format(strategy, var))

    ValueError: Trying to create optimizer slot variable under the scope for tf.distribute.Strategy (<tensorflow.python.distribute.one_device_strategy.OneDeviceStrategy object at 0x7efd64765bd0>), which is different from the scope used for the original variable (<tf.Variable 'tf_distil_bert_for_sequence_classification/distilbert/embeddings/tf_distil_bert_for_sequence_classification/distilbert/embeddings/word_embeddings/weight:0' shape=(28998, 768) dtype=float32, numpy=
    array([[-0.02513016, -0.03304445, -0.00243959, ..., -0.01084836,
            -0.04682418, -0.00948554],
           [-0.00482445, -0.02148623, -0.00871447, ..., -0.02602929,
            -0.03786189, -0.02410287],
           [-0.01653061, -0.01786226,  0.00105964, ..., -0.01637051,
            -0.03567044, -0.03141942],
           ...,
           [ 0.01190545, -0.02329331, -0.02250608, ..., -0.02713599,
            -0.04355597,  0.00010529],
           [ 0.00688736,  0.02267248,  0.02263871, ..., -0.00735895,
            -0.00814128,  0.00426289],
           [ 0.00320692, -0.0061747 ,  0.01624888, ...,  0.00641411,
             0.00060032,  0.01258053]], dtype=float32)>). Make sure the slot variables are created under the same strategy scope. This may happen if you're restoring from a checkpoint outside the scope
```

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
The same model trains successfully when trained as a tf.keras model with a batched tfdataset.
```
# modified training code to use the keras model instance that trains the model successfully.
def train_model(
    model_args, data_dir, model_dir, logs_dir, batch_size=32, num_epochs=10
):
    training_args = TFTrainingArguments(
        output_dir=model_dir,
        num_train_epochs=num_epochs,
        do_train=True,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size * 2,
        evaluation_strategy="steps",
        warmup_steps=500,
        weight_decay=0.01,
        logging_dir=logs_dir,
        dataloader_num_workers=15,
    )
    datasets = {
        "train": load_dataset(data_dir=data_dir, split="train", batch_size=batch_size),
        "val": load_dataset(
            data_dir=data_dir, split="validation", batch_size=batch_size
        ),
    }

    model = load_model(**model_args)

    history = model.fit(datasets["train"].batch(32), verbose=1)
    return history
```
Hello!

This is because you do not instantiate your model in the created strategy. You already have an example on how to train such models in the [repo](https://github.com/huggingface/transformers/tree/master/examples/text-classification)<|||||>Hi @jplu,
Are you referring to this: https://github.com/huggingface/transformers/blob/f33a6f34461fea61b579a7ec732fcd174b2b41cd/examples/text-classification/run_tf_text_classification.py#L263
i.e. do i just need to wrap `load_model` in the above code in the context manager `with training_args.strategy.scope():` resulting in 
```
# train model
def train_model(
    model_args, data_dir, model_dir, logs_dir, batch_size=32, num_epochs=10
):
    training_args = TFTrainingArguments(
        output_dir=model_dir,
        num_train_epochs=num_epochs,
        do_train=True,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size * 2,
        evaluation_strategy="steps",
        warmup_steps=500,
        weight_decay=0.01,
        logging_dir=logs_dir,
        dataloader_num_workers=15,
    )
    datasets = {
        "train": load_dataset(data_dir=data_dir, split="train", batch_size=batch_size),
        "val": load_dataset(
            data_dir=data_dir, split="validation", batch_size=batch_size
        ),
    }
    
    with training_args.strategy.scope():
        model = load_model(**model_args)

    trainer = TFTrainer(
        model=model,
        args=training_args,
        train_dataset=datasets["train"],
        eval_dataset=datasets["val"],
        compute_metrics=compute_metrics,
    )
    trainer.train()
    return trainer
```

If not could you point me to the right place?<|||||>Yes this is what I meant :)<|||||>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.