Make sure the slot variables are created under the same strategy scope.
Author: vlreinier
Pull request: no

- `transformers` version: 3.5.0
- Platform: jupyter notebook
- Python version: 3.6.9
- Tensorflow version (GPU?): 2.3.1
- Using GPU in script?: Single RTX 2080TI
- Model: "distilbert-base-multilingual-cased"

with tf.device('/device:GPU:0'):
    model.compile(optimizer=optimizer, loss=loss_fn)
    model.fit(train_dataset.batch(batch_size_train), epochs=1)

Code above works fine. Using TFtrainer with code below produces strategy error.
Note:
tf.config.list_physical_devices('GPU') gives [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
Steps to reproduce the behavior:

from transformers import BertTokenizer, TFBertForTokenClassification
from transformers import __version__
from transformers import TFTrainer, TFTrainingArguments

label2id = {"False": 0, "True": 1}
bert_model = "distilbert-base-multilingual-cased"
cache_dir = "cache/distilbert"

model = TFBertForTokenClassification.from_pretrained(
    bert_model,
    cache_dir = cache_dir, 
    num_labels = len(label2id), 
    label2id = label2id, 
    id2label = {v:k for k,v in label2id.items()}
)

training_args = TFTrainingArguments(
    output_dir='cp',   
    num_train_epochs=1,             
    per_device_train_batch_size=16, 
    warmup_steps=500,               
    weight_decay=0.05,              
)
                               
trainer = TFTrainer(
    model=model,                        
    args=training_args,                  
    train_dataset=train_dataset,    
)

trainer.train()

    ValueError: Trying to create optimizer slot variable under the scope for tf.distribute.Strategy (<tensorflow.python.distribute.one_device_strategy.OneDeviceStrategy object at 0x7effe7039588>), which is different from the scope used for the original variable (<tf.Variable 'tf_bert_for_token_classification/bert/embeddings/word_embeddings/weight:0' shape=(119547, 768) dtype=float32, numpy=
    array([[ 0.01447453, -0.03549159,  0.03377417, ..., -0.01235564,
             0.00624704, -0.01201372],
           [ 0.00689944,  0.00139387, -0.00716509, ...,  0.02127312,
            -0.00164859, -0.02350472],
           [ 0.00123599, -0.02220839, -0.01472212, ..., -0.02844208,
            -0.01958628,  0.01139562],
           ...,
           [ 0.03488934,  0.00115632,  0.0073231 , ...,  0.00768381,
            -0.02942067, -0.00667366],
           [ 0.00224815, -0.00895759,  0.0046453 , ...,  0.00379816,
             0.00176853, -0.01759749],
           [-0.02521203, -0.03274821, -0.00520367, ..., -0.01396327,
             0.0071948 , -0.01428833]], 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

Comments:
I think i already got an answer to my question from this page https://github.com/huggingface/transformers/blob/master/examples/token-classification/run_tf_ner.py :)

with training_args.strategy.scope():
    model = TFBertForTokenClassification.from_pretrained(bert_model,
                                                        cache_dir=cache_dir, 
                                                        num_labels=len(label2id), 
                                                        label2id=label2id, 
                                                        id2label={v:k for k,v in label2id.items()}
                                                        )
    model.summary()