fix dtype bug

[chujiezheng]

It would raise error if we use the bfloat16 dtype.

What does this PR do?

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[amyeroberts]

Hi @chujiezheng, thanks for opening a PR!

Could you share an example snippet which reproduces this error on main?

[chujiezheng]

import torch
from transformers import pipeline

rm_pipe = pipeline(
    "sentiment-analysis",
    model="weqweasdas/RM-Mistral-7B",
    device_map='auto',
    torch_dtype=torch.bfloat16,
)
texts = ['hello']
outputs = rm_pipe(texts)

[amyeroberts]

@chujiezheng Thanks for proving an example!

Could you also add a test which will fail on main but pass on this branch? Ideally, this test should be added to all pipelines to check that can accept bf16 inputs.

[chujiezheng]

Sorry, I do not get it well. What tests should I do, e.g., with other pipelines?