issue with --jit_mode_eval enabled in trainer commandline
### System Info


- `transformers` version: 4.24.0.dev0
- Platform: Linux-5.15.0-46-generic-x86_64-with-glibc2.31
- Python version: 3.9.13
- Huggingface_hub version: 0.10.1
- PyTorch version (GPU?): 1.12.1+cu102 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>


### Who can help?

Trainer: @sgugger

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

python3 examples/pytorch/text-classification/run_glue.py --model_name_or_path /skyrex01/wangyi/output/mrpc/ --task_name mrpc  --do_eval --max_seq_length 128 --output_dir /skyrex01/wangyi/output/mrpc/inference1/ --overwrite_output_dir True  --fp16 --jit_mode_eval



### Expected behavior

jit success and run the jit traced model.
current behavior
error log:
[INFO|trainer.py:557] 2022-10-30 20:10:11,309 >> Using cuda_amp half precision backend
10/30/2022 20:10:11 - INFO - __main__ - *** Evaluate ***
[INFO|trainer.py:725] 2022-10-30 20:10:11,309 >> The following columns in the evaluation set don't have a corresponding argument in `BertForSequenceClassification.forward` and have been ignored: idx, sentence2, sentence1. If idx, sentence2, sentence1 are not expected by `BertForSequenceClassification.forward`,  you can safely ignore this message.
Traceback (most recent call last):
  File "/home/wangyi/project/hugface/transformers/examples/pytorch/text-classification/run_glue.py", line 622, in <module>
    main()
  File "/home/wangyi/project/hugface/transformers/examples/pytorch/text-classification/run_glue.py", line 560, in main
    metrics = trainer.evaluate(eval_dataset=eval_dataset)
  File "/home/wangyi/project/hugface/transformers/src/transformers/trainer.py", line 2792, in evaluate
    output = eval_loop(
  File "/home/wangyi/project/hugface/transformers/src/transformers/trainer.py", line 2913, in evaluation_loop
    model = self._wrap_model(self.model, training=False, dataloader=dataloader)
  File "/home/wangyi/project/hugface/transformers/src/transformers/trainer.py", line 1299, in _wrap_model
    model = self.torch_jit_model_eval(model, dataloader, training)
  File "/home/wangyi/project/hugface/transformers/src/transformers/trainer.py", line 1263, in torch_jit_model_eval
    jit_model = torch.jit.trace(jit_model, jit_inputs, strict=False)
  File "/skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/jit/_trace.py", line 750, in trace
    return trace_module(
  File "/skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/jit/_trace.py", line 967, in trace_module
    module._c._create_method_from_trace(
  File "/skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1118, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/wangyi/project/hugface/transformers/src/transformers/models/bert/modeling_bert.py", line 1552, in forward
    outputs = self.bert(
  File "/skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1118, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/wangyi/project/hugface/transformers/src/transformers/models/bert/modeling_bert.py", line 968, in forward
    batch_size, seq_length = input_shape
ValueError: not enough values to unpack (expected 2, got 1)
<|||||>case2: if we run only predict. command like
python3 examples/pytorch/text-classification/run_glue.py --model_name_or_path /skyrex01/wangyi/output/mrpc/ --task_name mrpc  --do_predict --max_seq_length 128 --output_dir /skyrex01/wangyi/output/mrpc/inference1/ --overwrite_output_dir True  --fp16 --jit_mode_eval

error changing, since inputdata does not contain "labels" in this case
jit failure as "failed to use PyTorch jit mode due to: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select)."
jit error:
[INFO|modeling_utils.py:2616] 2022-10-30 20:13:22,055 >> All the weights of BertForSequenceClassification were initialized from the model checkpoint at /skyrex01/wangyi/output/mrpc/.
If your task is similar to the task the model of the checkpoint was trained on, you can already use BertForSequenceClassification for predictions without further training.
10/30/2022 20:13:22 - WARNING - datasets.arrow_dataset - Loading cached processed dataset at /skyrex01/wangyi/.cache/datasets/glue/mrpc/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad/cache-1928edc8ebbd0881.arrow
10/30/2022 20:13:22 - WARNING - datasets.arrow_dataset - Loading cached processed dataset at /skyrex01/wangyi/.cache/datasets/glue/mrpc/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad/cache-3447598ff1e90f2d.arrow
Running tokenizer on dataset:   0%|                                                                                                                                         | 0/2 [00:00<?, ?ba/s]
10/30/2022 20:13:22 - INFO - datasets.arrow_dataset - Caching processed dataset at /skyrex01/wangyi/.cache/datasets/glue/mrpc/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad/cache-64cad0d6db19155f.arrow
Running tokenizer on dataset:  50%|████████████████████████████████████████████████████████████████▌                                                                | 1/2 [00:00<00:00,  7.20ba/s]
[INFO|trainer.py:557] 2022-10-30 20:13:25,460 >> Using cuda_amp half precision backend
10/30/2022 20:13:25 - INFO - __main__ - *** Predict ***
[INFO|trainer.py:725] 2022-10-30 20:13:25,462 >> The following columns in the test set don't have a corresponding argument in `BertForSequenceClassification.forward` and have been ignored: idx, sentence1, sentence2. If idx, sentence1, sentence2 are not expected by `BertForSequenceClassification.forward`,  you can safely ignore this message.
[WARNING|trainer.py:1268] 2022-10-30 20:13:25,714 >> failed to use PyTorch jit mode due to: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select).
[INFO|trainer.py:2925] 2022-10-30 20:13:25,715 >> ***** Running Prediction *****
[INFO|trainer.py:2927] 2022-10-30 20:13:25,715 >>   Num examples = 1725
[INFO|trainer.py:2930] 2022-10-30 20:13:25,715 >>   Batch size = 8
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 216/216 [00:02<00:00, 87.92it/s]
10/30/2022 20:13:28 - INFO - __main__ - ***** Predict results mrpc *****
[INFO|modelcard.py:444] 2022-10-30 20:13:28,370 >> Dropping the following result as it does not have all the necessary fields:
{'task': {'name': 'Text Classification', 'type': 'text-classification'}, 'dataset': {'name': 'GLUE MRPC', 'type': 'glue', 'args': 'mrpc'}}

<|||||>case3: if we run only predict on cpu. command like
python3 examples/pytorch/text-classification/run_glue.py --model_name_or_path /skyrex01/wangyi/output/mrpc/ --task_name mrpc  --do_predict --max_seq_length 128 --output_dir /skyrex01/wangyi/output/mrpc/inference1/ --overwrite_output_dir True  --bf16 --jit_mode_eval --no_cuda
error pop like
ERROR: Tensor-valued Constant nodes differed in value across invocations. This often indicates that the tracer has encountered untraceable code.
        Node:
                %13 : Tensor = prim::Constant[value={8}](), scope: __module.bert/__module.bert.encoder/__module.bert.encoder.layer.0/__module.bert.encoder.layer.0.attention/__module.bert.encoder.layer.0.attention.self # /home/wangyi/project/hugface/transformers/src/transformers/models/bert/modeling_bert.py:341:0
        Source Location:
                /home/wangyi/project/hugface/transformers/src/transformers/models/bert/modeling_bert.py(341): forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1118): _slow_forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1130): _call_impl
                /home/wangyi/project/hugface/transformers/src/transformers/models/bert/modeling_bert.py(419): forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1118): _slow_forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1130): _call_impl
                /home/wangyi/project/hugface/transformers/src/transformers/models/bert/modeling_bert.py(489): forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1118): _slow_forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1130): _call_impl
                /home/wangyi/project/hugface/transformers/src/transformers/models/bert/modeling_bert.py(603): forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1118): _slow_forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1130): _call_impl
                /home/wangyi/project/hugface/transformers/src/transformers/models/bert/modeling_bert.py(1014): forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1118): _slow_forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1130): _call_impl
                /home/wangyi/project/hugface/transformers/src/transformers/models/bert/modeling_bert.py(1552): forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1118): _slow_forward
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/nn/modules/module.py(1130): _call_impl
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/jit/_trace.py(967): trace_module
                /skyrex05/wangyi/miniconda3/envs/compatibility_test/lib/python3.9/site-packages/torch/jit/_trace.py(750): trace
                /home/wangyi/project/hugface/transformers/src/transformers/trainer.py(1263): torch_jit_model_eval
                /home/wangyi/project/hugface/transformers/src/transformers/trainer.py(1299): _wrap_model
                /home/wangyi/project/hugface/transformers/src/transformers/trainer.py(2913): evaluation_loop
                /home/wangyi/project/hugface/transformers/src/transformers/trainer.py(2866): predict
                /home/wangyi/project/hugface/transformers/examples/pytorch/text-classification/run_glue.py(588): main
                /home/wangyi/project/hugface/transformers/examples/pytorch/text-classification/run_glue.py(622): <module>
        Comparison exception:   The values for attribute 'shape' do not match: torch.Size([]) != torch.Size([768, 768]).
<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.