🐛 [TF_BART] "<internal expr>" has dtype float32 in the TRUE branch, but dtype=int32 in the FALSE branch
Author: astariul
Pull request: no

## 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.0.dev0
- Platform: Linux-4.19.0-13-cloud-amd64-x86_64-with-debian-10.7
- Python version: 3.7.3
- PyTorch version (GPU?): 1.7.0 (False)
- Tensorflow version (GPU?): 2.2.0 (False)
- Using GPU in script?: No (TPU)
- Using distributed or parallel set-up in script?: Yes

### Who can help
 
@patrickvonplaten

## Information

Model I am using (Bert, XLNet ...): TFBart

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

## To reproduce

When I try to run TF_Bart on TPU, I'm getting the following error :

> TypeError: "<internal expr>" has dtype float32 in the TRUE branch, but dtype=int32 in the FALSE branch. TensorFlow control flow requires that they are the same.

It seems to come from the dropout operation :

https://github.com/huggingface/transformers/blob/b01ddc9577b87f057e163d49563ee3f74f4810cf/src/transformers/models/bart/modeling_tf_bart.py#L373

<details>
 <summary> Full stack trace (click to expand...)</summary>

>2020/12/11 00:00:55 - INFO - transformers_addons.trainer_tf -   ***** Running Evaluation *****
2020/12/11 00:00:55 - INFO - transformers_addons.trainer_tf -     Batch size = 8
The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
Traceback (most recent call last):
  File "train.py", line 203, in <module>
    main()
  File "train.py", line 194, in main
    result = trainer.evaluate()
  File "/home/remondnicola/text-summarization/transformers_addons/trainer_tf.py", line 281, in evaluate
    output = self._prediction_loop(eval_dataset, description="Evaluation", prediction_loss_only=prediction_loss_only)
  File "/home/remondnicola/text-summarization/transformers_addons/trainer_tf.py", line 207, in _prediction_loop
    loss, logits = self._evaluate_steps(features, labels)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 580, in __call__
    result = self._call(*args, **kwds)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 627, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 506, in _initialize
    *args, **kwds))
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 2446, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 2777, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 2667, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 981, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 441, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 3299, in bound_method_wrapper
    return wrapped_fn(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 968, in wrapper
    raise e.ag_error_metadata.to_exception(e)
TypeError: in user code:
>
>    /home/remondnicola/text-summarization/transformers_addons/trainer_tf.py:169 _evaluate_steps  *
        per_replica_loss, per_replica_logits = self.args.strategy.experimental_run_v2(
    train.py:29 _run_model  *
        out = self.model(features, training=training, **labels)
    /home/remondnicola/text-summarization/transformers_addons/models/bart/modeling_tf_bart.py:88 call  *
        outputs = super().call(inputs["input_ids"],
    /home/remondnicola/.venv/summarization/lib/python3.7/site-packages/transformers/models/bart/modeling_tf_bart.py:1110 call  *
        outputs = self.model(
    /home/remondnicola/.venv/summarization/lib/python3.7/site-packages/transformers/models/bart/modeling_tf_bart.py:977 call  *
        inputs["encoder_outputs"] = self.encoder(
    /home/remondnicola/.venv/summarization/lib/python3.7/site-packages/transformers/models/bart/modeling_tf_bart.py:373 call  *
        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py:924 if_stmt
        basic_symbol_names, composite_symbol_names)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py:962 tf_if_stmt
        error_checking_orelse)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/deprecation.py:507 new_func
        return func(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/control_flow_ops.py:1177 cond
        return cond_v2.cond_v2(pred, true_fn, false_fn, name)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/cond_v2.py:91 cond_v2
        op_return_value=pred)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py:981 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py:958 error_checking_orelse
        basic_symbol_names + composite_symbol_names)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py:298 _verify_tf_cond_vars
        functools.partial(_verify_single_cond_var, name), body_var, orelse_var)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/nest.py:617 map_structure
        structure[0], [func(*x) for x in entries],
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/nest.py:617 <listcomp>
        structure[0], [func(*x) for x in entries],
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/operators/control_flow.py:267 _verify_single_cond_var
        orelse_var.dtype.name))
>
>    TypeError: "<internal expr>" has dtype float32 in the TRUE branch, but dtype=int32 in the FALSE branch. TensorFlow control flow requires that they are the same.

</details>

Comments: