DeformableDETR two stage not support bfloat16

[DonggeunYu]

System Info

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

• transformers version: 4.39.0
• Platform: Linux-5.4.0-81-generic-x86_64-with-glibc2.35
• Python version: 3.10.12
• Huggingface_hub version: 0.19.4
• Safetensors version: 0.4.2
• Accelerate version: 0.27.2
• Accelerate config: not found
• PyTorch version (GPU?): 2.1.0+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 GPU in script?:
• Using distributed or parallel set-up in script?:

Who can help?

@amyeroberts

Information

• The official example scripts
• 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

import torch
import requests
from PIL import Image
from transformers import AutoImageProcessor
from transformers.models.deformable_detr import DeformableDetrForObjectDetection


url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

processor = AutoImageProcessor.from_pretrained("SenseTime/deformable-detr-with-box-refine-two-stage")
model = DeformableDetrForObjectDetection.from_pretrained(
    "SenseTime/deformable-detr-with-box-refine-two-stage",
)
print(model.config.disable_custom_kernels)
model = model.to(torch.bfloat16)
model = model.to('cuda')
model = model.eval()
inputs = processor(images=image, return_tensors="pt")

inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16).to('cuda')
inputs['pixel_mask'] = inputs['pixel_mask'].to('cuda')

with torch.no_grad():
    outputs = model(**inputs)
print(outputs)

Expected behavior

Traceback (most recent call last):
  File "/nas/k8s/dev/mlops/donggeun/tmp/selfsup/vision_language/4.py", line 24, in <module>
    outputs = model(**inputs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/deformable_detr/modeling_deformable_detr.py", line 1976, in forward
    outputs = self.model(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/deformable_detr/modeling_deformable_detr.py", line 1812, in forward
    pos_trans_out = self.pos_trans_norm(self.pos_trans(topk_coords_logits))
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: mat1 and mat2 must have the same dtype, but got BFloat16 and Float