Tensor size mismatch when trying to run RT-DETR on multiple gpus

### System Info - `transformers` version: 4.44.2 - Platform: Linux-5.4.0-174-generic-x86_64-with-glibc2.31 - Python version: 3.11.6 - Huggingface_hub version: 0.24.6 - Safetensors version: 0.4.4 - Accelerate version: 0.33.0 - Accelerate config: not found - PyTorch version (GPU?): 2.1.2+cu121 (True) - Tensorflow version (GPU?): not installed (NA) - Flax version (CPU?/GPU?/TPU?): not installed (NA) - Jax version: not installed - JaxLib version: not installed - Using distributed or parallel set-up in script?: <fill in> - Using GPU in script?: <fill in> - GPU type: Tesla V100-DGXS-16GB ### Who can help? @amyeroberts @muellerz @SunMarc ### Information - [X] The official example scripts - [X] My own modified scripts ### Tasks - [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...) - [ ] My own task or dataset (give details below) ### Reproduction Following the example on [the official pytorch example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/object-detection/run_object_detection.py) and [here](https://huggingface.co/docs/transformers/en/tasks/object_detection) it seems that I get the stack trace below after following these steps: 1. Set the initial model class to RT-DETR assuming other parts of the example have been followed ``` IMAGE_SIZE = 1280 CHECKPOINT = "PekingU/rtdetr_r50vd_coco_o365" DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = AutoModelForObjectDetection.from_pretrained( CHECKPOINT, id2label=id2label, label2id=label2id, anchor_image_size=None, ignore_mismatched_sizes=True ) ``` 2. Set the batch size to 4 and set to 4 visible GPUs assuming other parts of the example have been followed ``` os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" training_args = TrainingArguments( output_dir=output_path, num_train_epochs=20, max_grad_norm=0.1, learning_rate=5e-5, warmup_steps=300, per_device_train_batch_size=4, dataloader_num_workers=…