Errors while using TFAutoModelForMultipleChoice and TFTrainer on winogrande dataset 
# ❓ Questions & Help
I'm using TFAutoModelForMultipleChoice and TFTrainer to train on winogrande dataset.
The dataset is from [here](https://leaderboard.allenai.org/winogrande/submissions/get-started)

And I got an error while I starting training.
```python
Traceback (most recent call last):
  File "run_wsc.py", line 233, in <module>
    main()
  File "run_wsc.py", line 208, in main
    trainer.train()
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/transformers/trainer_tf.py", line 358, in train
    for step, training_loss in enumerate(self._training_steps(train_ds, optimizer)):
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/transformers/trainer_tf.py", line 401, in _training_steps
    for i, loss in enumerate(self._accumulate_next_gradients(ds)):
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/transformers/trainer_tf.py", line 434, in _accumulate_next_gradients
    yield _accumulate_next()
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py", line 568, in __call__
    result = self._call(*args, **kwds)
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py", line 615, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py", line 497, in _initialize
    *args, **kwds))
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 2389, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 2703, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 2593, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.py", line 978, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py", line 439, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.py", line 968, in wrapper
    raise e.ag_error_metadata.to_exception(e)
TypeError: in converted code:

    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/transformers/trainer_tf.py:430 _accumulate_next  *
        return self._accumulate_gradients(per_replica_features, per_replica_labels)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/transformers/trainer_tf.py:440 _accumulate_gradients  *
        per_replica_loss = self.args.strategy.experimental_run_v2(
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/distribute/one_device_strategy.py:180 experimental_run_v2
        return super(OneDeviceStrategy, self).experimental_run_v2(fn, args, kwargs)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/transformers/trainer_tf.py:453 _forward  *
        per_example_loss, _ = self._run_model(features, labels, True)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/transformers/trainer_tf.py:474 _run_model  *
        loss, logits = self.model(features, labels=labels, training=training)[:2]
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py:778 __call__
        outputs = call_fn(cast_inputs, *args, **kwargs)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/transformers/modeling_tf_bert.py:1114 call  *
        loss = self.compute_loss(labels, reshaped_logits)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/transformers/modeling_tf_utils.py:134 compute_loss  *
        if shape_list(logits)[1] == 1:
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:918 if_stmt
        basic_symbol_names, composite_symbol_names)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:956 tf_if_stmt
        error_checking_orelse)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/util/deprecation.py:507 new_func
        return func(*args, **kwargs)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py:1174 cond
        return cond_v2.cond_v2(pred, true_fn, false_fn, name)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/ops/cond_v2.py:83 cond_v2
        op_return_value=pred)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.py:983 func_graph_from_py_func
        expand_composites=True)
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/util/nest.py:568 map_structure
        structure[0], [func(*x) for x in entries],
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/util/nest.py:568 <listcomp>
        structure[0], [func(*x) for x in entries],
    /Users/qixinli/anaconda/envs/a2/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.py:943 convert
        (str(python_func), type(x)))

    TypeError: To be compatible with tf.contrib.eager.defun, Python functions must return zero or more Tensors; in compilation of <function tf_if_stmt.<locals>.error_checking_body at 0x158e993b0>, found return value of type <class 'tensorflow.python.keras.losses.MeanSquaredError'>, which is not a Tensor.
```

## Details

```python
import logging
import os
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, Optional

import numpy as np
import tensorflow as tf

from transformers import (
    AutoConfig,
    AutoTokenizer,
    EvalPrediction,
    HfArgumentParser,
    PreTrainedTokenizer,
    TFTrainer,
    TFTrainingArguments,
    TFAutoModelForMultipleChoice
)

from wsc_processor import (
    WinograndeProcessor,
    convert_multiple_choice_examples_to_features,
    compute_metrics
)

class Split(Enum):
    train = "train"
    dev = "validation"
    test = "test"

def get_wsc_ds(
    tokenizer: PreTrainedTokenizer, max_seq_length: Optional[int] = None, mode: Split = Split.train, data_dir:str=None
):
    processor = WinograndeProcessor()
    if mode == Split.train:
        examples = processor.get_train_examples(data_dir)
    elif mode == Split.dev:
        examples = processor.get_dev_examples(data_dir)
    else:
        examples = processor.get_test_examples(data_dir)

    features = convert_multiple_choice_examples_to_features(examples, label_list=processor.get_labels(), tokenizer=tokenizer, max_seq_length=max_seq_length, output_mode="multiple_choice")

    def gen():
        for ex in features:
            inputs = []
            masks = []
            segments = []
            for feature in ex.option_features:
                inputs.append(feature["input_ids"])
                masks.append(feature["input_mask"])
                segments.append(feature["segment_ids"])
            yield (
                {
                    "input_ids": inputs,
                    "attention_mask": masks,
                    "token_type_ids": segments,
                },
                ex.label,
            )

    ds = tf.data.Dataset.from_generator(
        gen,
        ({"input_ids": tf.int32, "attention_mask": tf.int32, "token_type_ids": tf.int32}, tf.int64),
        (
            {
                "input_ids": tf.TensorShape([None,None]),
                "attention_mask": tf.TensorShape([None,None]),
                "token_type_ids": tf.TensorShape([None,None]),
            },
            tf.TensorShape([]),
        ),
    )
    return ds

logger = logging.getLogger(__name__)

@dataclass
class GlueDataTrainingArguments:
    task_name: str = field(default="wsc")
    data_dir: str = field(default="./data/WSC")
    max_seq_length: int = field(
        default=128,
        metadata={
            "help": "The maximum total input sequence length after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded."
        },
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )

@dataclass
class ModelArguments:
    model_name_or_path: str = field(
        default="roberta-base",
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    use_fast: bool = field(default=False, metadata={"help": "Set this flag to use fast tokenization."})
    cache_dir: Optional[str] = field(
        default="./model/", metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )

def main():
    parser = HfArgumentParser((ModelArguments, GlueDataTrainingArguments, TFTrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if (
        os.path.exists(training_args.output_dir)
        and os.listdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
        )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )

    logger.info("Training/evaluation parameters %s", training_args)

    output_mode = "classification"

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        # num_labels=1,
        finetuning_task=data_args.task_name,
        cache_dir=model_args.cache_dir
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir
    )

    # with training_args.strategy.scope():
    model = TFAutoModelForMultipleChoice.from_pretrained(
        model_args.model_name_or_path,
        from_pt=bool(".bin" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
    )

    # Get datasets
    train_dataset = (
        get_wsc_ds(tokenizer=tokenizer, max_seq_length=data_args.max_seq_length,data_dir=data_args.data_dir)
        if training_args.do_train
        else None
    )
    eval_dataset = (
        get_wsc_ds(tokenizer=tokenizer, max_seq_length=data_args.max_seq_length,mode=Split.dev, data_dir=data_args.data_dir)
        if training_args.do_eval
        else None
    )

    def metric(p: EvalPrediction) -> Dict:
        if output_mode == "classification":
            preds = np.argmax(p.predictions, axis=1)
        else:
            preds = np.squeeze(p.predictions)
        return compute_metrics("winogrande", preds, p.label_ids)

    # Initialize our Trainer
    trainer = TFTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=metric,
    )

    # Training
    if training_args.do_train:
        trainer.train()
        trainer.save_model()
        tokenizer.save_pretrained(training_args.output_dir)
    return results


if __name__ == "__main__":
    main()
```
the wsc_processor is from [here](https://github.com/allenai/winogrande/blob/master/scripts/utils.py)

I wonder what is wrong in the code (am I wrong in the get_wsc_ds() method ?) , and how to fix it?

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.
