Runtime error when training DetForObjectDetection using HFTrainer with GPU.
Author: jnishi
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.9.2
- Platform: Linux-5.4.0-81-generic-x86_64-with-glibc2.27
- Python version: 3.8.0
- PyTorch version (GPU?): 1.9.0+cu111 (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?: <yes>
- Using distributed or parallel set-up in script?: <no>

## Information

Model I am using: DetrForObjectDetection

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

I'm training DetrForObjectDetection by using HFTrainer. 

Save a script file below as `mini_example.py`, and run as `python  mini_example.py --output_dir mini_model mini_model` after setting `img_folder` to the path to the coco image dataset folder and  `annotations` to the path to the coco annotation JSON file.

```python
from typing import Dict, List, Union

import torch
from torchvision.datasets import CocoDetection

from transformers import (
    DetrConfig,
    DetrFeatureExtractor,
    DetrForObjectDetection,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
)


def load_category(category):
    id2label = {}
    label2id = {}
    maxid = 0
    for k, v in category.items():
        id2label[int(k)] = v["name"]
        label2id[v["name"]] = int(k)
        maxid = max(maxid, int(k))
    for i in range(maxid):
        if not (i in id2label):
            id2label[i] = None
    return id2label, label2id


class DetrData(CocoDetection):
    def __init__(self, img_folder, annotations, feature_extractor, train=True):
        super(DetrData, self).__init__(img_folder, annotations)
        self.feature_extractor = feature_extractor

    def __getitem__(self, idx):
        # read in PIL image and target in COCO format                                                                                                                                                                                                                                                       img, target = super(DetrData, self).__getitem__(idx)

        # preprocess image and target (converting target to DETR format, resizing + normalization of both image and target)                                                                                                                                                                                 image_id = self.ids[idx]
        target = {'image_id': image_id, 'annotations': target}
        encoding = self.feature_extractor(images=img, annotations=target, return_tensors="pt")
        encoding["pixel_values"] = encoding["pixel_values"].squeeze() # remove batch dimension                                                                                                                                                                                                              encoding["labels"] = encoding["labels"][0] # remove batch dimension                                                                                                                                                                                                                                 return encoding

@dataclass
class DataCollatorDetr:

    feature_extractor: DetrFeatureExtractor

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        pixel_values = [item["pixel_values"] for item in features]
        encoding = self.feature_extractor.pad_and_create_pixel_mask(pixel_values, return_tensors="pt")
        encoding["labels"] = [item["labels"] for item in features]
        return encoding


def main():
    parser = HfArgumentParser((TrainingArguments))
    training_args, = parser.parse_args_into_dataclasses()
    feature_extractor = DetrFeatureExtractor()
    train_dataset = DetrData(img_folder="path/to/image_folder", annotations="path/to/annotation_file", feature_extractor=feature_extractor)                                                                                                                                                               
    id2label, label2id = load_category(train_dataset.coco.cats)
    config = DetrConfig.from_pretrained("facebook/detr-resnet-50")
    config.id2label = id2label
    config.label2id = label2id

    model = DetrForObjectDetection.from_pretrained(
        "facebook/detr-resnet-50",
        config=config)

    # Initialize our Trainer                                                                                                                                                                                                                                                                            
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        tokenizer=feature_extractor,
        data_collator=DataCollatorDetr(feature_extractor=feature_extractor),
    )

    train_result = trainer.train()

if __name__ == "__main__":
    main()
```

When train without GPU, it works fine, but got RuntimeError below with GPU,

```
Traceback (most recent call last):
  File "mini_example.py", line 97, in <module>
    main()
  File "mini_example.py", line 93, in main
    train_result = trainer.train()
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/transformers/trainer.py", line 1286, in train
    tr_loss += self.training_step(model, inputs)
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/transformers/trainer.py", line 1779, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/transformers/trainer.py", line 1811, in compute_loss
    outputs = model(**inputs)
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/transformers/models/detr/modeling_detr.py", line 1435, in forward
    loss_dict = criterion(outputs_loss, labels)
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/transformers/models/detr/modeling_detr.py", line 2009, in forward
    indices = self.matcher(outputs_without_aux, targets)
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/transformers/models/detr/modeling_detr.py", line 2127, in forward
    bbox_cost = torch.cdist(out_bbox, tgt_bbox, p=1)
  File "/home/christopher/pyenvs/detr/lib/python3.8/site-packages/torch/functional.py", line 1049, in cdist
    return _VF.cdist(x1, x2, p, None)  # type: ignore[attr-defined]
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking arugment for argument x2 in method wrapper__cdist_forward)
  0%|                                                                                                                                                                                                                                                                      | 0/1875 [00:03<?, ?it/s]
  ```

This is maybe because `inputs["labels"]` is not sent to GPU here https://github.com/huggingface/transformers/blob/v4.9.2/src/transformers/trainer.py#L1734 which is called at 
https://github.com/huggingface/transformers/blob/master/src/transformers/trainer.py#L1771 because it is dict. 

Any suggestion on how to fix it?

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior
Successfully complete training
<!-- A clear and concise description of what you would expect to happen. -->


Comments:
Hey @jnishi,

Thanks a lot for your issue!
Could you please try to make a minimum reproducible code example that doesn't force us to manually create a `img_folder` or `annotations` folder? Ideally, you could link to a colab that runs in less than a minute to reproduce the error. 

Also cc'ing @NielsRogge here for DETR

Here is the link to colab.
https://colab.research.google.com/drive/1qvasKfJGhxoNn-l_5GZwkvh4FhW59gBS?usp=sharing

Please upload sample.jpg and sample.json included below before you run colab.
[detr_samples.tar.gz](https://github.com/huggingface/transformers/files/7011436/detr_samples.tar.gz)


Thanks for the colab! It was indeed easy to reproduce the issue.

I've fixed it here: https://colab.research.google.com/drive/1oIHGwr1U0sw-6KW-MG60s-ksXA-kYyUO?usp=sharing 

As you already spotted, the problem is in the `_prepare_inputs()` method of the Trainer, which does not take into account inputs which are lists. For DETR, the `labels` are a list of dictionaries, each dictionary containing the annotations (class labels and boxes) for an example in the batch. I've fixed it by overwriting that method.

cc'ing @sgugger, as this could be incorporated directly in the Trainer, instead of having to overwrite it.

Thanks for a quick response, and suggestion of the fix. It works fine in my scripts too.
I would be more than happy to incorporate it directly.

BTW, I have another problem with a multi-GPU environment, so I created another issue.
https://github.com/huggingface/transformers/issues/13197

The PR linked above should solve this problem. It's a bit more general than your solution in the notebook @NielsRogge to handle any nested dict/list of tensors.