Errors while using TFAutoModelForMultipleChoice and TFTrainer on winogrande dataset 
Author: QixinLi
Pull request: no
Labels: name: wontfix

# ❓ 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?


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.