TFRobertaForSequenceClassification fails on TPU on Transformers >2.0.0
Author: rickysaurav
Pull request: no
Labels: name: wontfix

## 🐛 Bug

<!-- Important information -->

Model I am using (TFRobertaForSequenceClassification):

Language I am using the model on (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:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details)

## To Reproduce

Steps to reproduce the behavior:

1. Use a TPU runtime on colab 
2. ```python
    resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
    tf.config.experimental_connect_to_cluster(resolver)
    tf.tpu.experimental.initialize_tpu_system(resolver)
    strategy = tf.distribute.experimental.TPUStrategy(resolver)
    with tf.device('/job:worker'):
      with strategy.scope():
        # model = TFRobertaForSequenceClassification.from_pretrained('bert-large-uncased-whole-word-masking',num_labels = (len(le.classes_)))
        model = TFRobertaForSequenceClassification.from_pretrained('roberta-large',num_labels = 2)
        print('model loaded')
        inp  = np.random.randint(10,100, size=(12800, 64))
        inp[:,0]=0
        inp[:,63]=2
        labels = np.random.randint(2,size = (12800,1))
        print('data generated')
        optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0)
        loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
        metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
        model.compile(optimizer=optimizer, loss=loss, metrics=[metric])
        print('starting fitting')
        # model.fit([train_input_ids,train_input_masks],y_train,epochs = 3,batch_size = 64,validation_data=([test_input_ids,test_input_masks], y_test),verbose=1)
        model.fit(inp,labels,epochs = 2,batch_size = 64,verbose=1)



## Environment

* OS: Google Colab TPU runtime
* Python version:3.6
* PyTorch version:NA
* PyTorch Transformers version (or branch):2.1.1
* Using GPU ? No
* Distributed of parallel setup ? Yes

## Additional context

The following error gets thrown when calling model.fit() 
```

InvalidArgumentError                      Traceback (most recent call last)

<ipython-input-4-b77065bb89ae> in <module>()
     15       print('starting fitting')
     16       # model.fit([train_input_ids,train_input_masks],y_train,epochs = 3,batch_size = 64,validation_data=([test_input_ids,test_input_masks], y_test),verbose=1)
---> 17       model.fit(inp,labels,epochs = 2,batch_size = 64,verbose=1)

11 frames

/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

InvalidArgumentError:  Compilation failure: Detected unsupported operations when trying to compile graph tf_roberta_for_sequence_classification_roberta_cond_true_122339[] on XLA_TPU_JIT: PrintV2 (No registered 'PrintV2' OpKernel for XLA_TPU_JIT devices compatible with node {{node PrintV2}}
	.  Registered:  device='CPU'
){{node PrintV2}}
	 [[tf_roberta_for_sequence_classification/roberta/cond]]
	TPU compilation failed
	 [[tpu_compile_succeeded_assert/_5504150486074133972/_3]]
Additional GRPC error information:
{"created":"@1571518085.015232162","description":"Error received from peer","file":"external/grpc/src/core/lib/surface/call.cc","file_line":1039,"grpc_message":" Compilation failure: Detected unsupported operations when trying to compile graph tf_roberta_for_sequence_classification_roberta_cond_true_122339[] on XLA_TPU_JIT: PrintV2 (No registered 'PrintV2' OpKernel for XLA_TPU_JIT devices compatible with node {{node PrintV2}}\n\t.  Registered:  device='CPU'\n){{node PrintV2}}\n\t [[tf_roberta_for_sequence_classification/roberta/cond]]\n\tTPU compilation failed\n\t [[tpu_compile_succeeded_assert/_5504150486074133972/_3]]","grpc_status":3} [Op:__inference_distributed_function_154989]

Function call stack:
distributed_function -> distributed_function
```
The reason behind this error seems to be the tf.print() in the following code , which is not supported on TPU. 

https://github.com/huggingface/transformers/blob/82f6abd98aaa691ca0adfe21e85a17dc6f386497/transformers/modeling_tf_roberta.py#L78-L80


Comments:
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.