Conditional branching logic in modeling_tf_flaubert.py causing errors with TF Graph
Author: ANarayan
Pull request: no

Hi @jplu !
I am encountering an error when running the TFFlaubert model inside of a tensorflow graph.

Here is some code to reproduce the issue:

```
import FlaubertTokenizer, TFFlaubertModel, FlaubertConfig
import tensorflow as tf

config=FlaubertConfig.from_pretrained('jplu/tf-flaubert-small-cased', output_attentions=True, output_hidden_states=True, return_dict=True)
tokenizer = TFFlaubertModel.from_pretrained(config=config, pretrained_model_name_or_path='jplu/tf-flaubert-small-cased')

@tf.function
def train_step(inputs, mask, token_type_ids):
    with tf.GradientTape() as tape:
        a = model({
            "input_ids": inputs,
            "training": True,
            "attention_mask": mask,
            "token_type_ids": token_type_ids,
        })
        
train_step(inputs, mask, token_type_ids)
```

The error seems to be caused by L611-624 in modeling_tf_flaubert.py [here](https://github.com/huggingface/transformers/blob/c60e0e1ee45f4bf1017736b146c51729f120bb83/src/transformers/models/flaubert/modeling_tf_flaubert.py#L611)

The error message is as follows:

> TypeError: in user code:
> 
>    python-input-5-4a1e131ff478:4 train_step  *
>         a = model({
>     /Users/ludwig/venv/lib/python3.6/site-packages/transformers/models/flaubert/modeling_tf_flaubert.py:274 call  *
>         outputs = self.transformer(
>     /Users/ludwig/venv/lib/python3.6/site-packages/transformers/models/flaubert/modeling_tf_flaubert.py:616 call  *
>         for i in range(self.n_layers):
>     /Users/ludwig/venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py:1163 if_stmt
>         _tf_if_stmt(cond, body, orelse, get_state, set_state, symbol_names, nouts)
>     /Users/ludwig/venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py:1210 _tf_if_stmt
>         cond, aug_body, aug_orelse, strict=True)
>     /Users/udwig/venv/lib/python3.6/site-packages/tensorflow/python/util/dispatch.py:201 wrapper
>         return target(*args, **kwargs)
>     /Users/ludwig/venv/lib/python3.6/site-packages/tensorflow/python/util/deprecation.py:538 new_func
>         return func(*args, **kwargs)
>     /Users/ludwig/venv/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py:1180 cond
>         return cond_v2.cond_v2(pred, true_fn, false_fn, name)
>     /Users/ludwig/venv/lib/python3.6/site-packages/tensorflow/python/ops/cond_v2.py:96 cond_v2
>         op_return_value=pred)
>     /Users/ludwig/venv/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py:990 func_graph_from_py_func
>         func_outputs = python_func(*func_args, **func_kwargs)
>     /Users/ludwig/venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py:1206 aug_orelse
>         _verify_tf_cond_vars(new_body_vars_[0], new_orelse_vars, symbol_names)
>     /Users/ludwig/venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py:365 _verify_tf_cond_vars
>         ' branches:\n\n{}'.format(name, str(e)))
> 
>     TypeError: 'hidden_states' must have the same nested structure in the main and else branches:
>     
>     The two structures don't have the same nested structure.
>     
>     First structure: type=tuple str=(<tf.Tensor 'tf_flaubert_model/transformer/mul:0' shape=(18, 44, 512) dtype=float32>,)
>     
>     Second structure: type=tuple str=()
>     
>     More specifically: The two structures don't have the same number of elements. First structure: type=tuple str=(<tf.Tensor 'tf_flaubert_model/transformer/mul:0' shape=(18, 44, 512) dtype=float32>,). Second structure: type=tuple str=()
>     Entire first structure:
>     (.,)
>     Entire second structure:
>     ()


Comments:
Hello,

To help you we need to have more information such as your environement info, and a standalone piece of code to let us reproduce your error. Thanks!

As a first guess, I can say that the issue you get might come from a misformed input, can you try with:
```
@tf.function
def train_step(inputs, mask, token_type_ids):
    with tf.GradientTape() as tape:
        a = model({
            "input_ids": inputs
            "attention_mask": mask,
            "token_type_ids": token_type_ids,
        }, training=True)
```

That worked! Thank you!!