blip2 RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.HalfTensor) should be

源代码

from transformers import AutoProcessor, Blip2ForConditionalGeneration
import torch

model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16)  # load in float16

inputs = processor(image, text=prompt, return_tensors="pt").to(device, torch.float16)

generated_ids = model.generate(**inputs, max_new_tokens=20)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()

报错信息

File "lib/python3.9/site-packages/torch/nn/modules/conv.py", line 442, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.HalfTensor) should be the same

 解决方案

在导入model时增加device_map参数

model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", device_map = "auto", torch_dtype=torch.float16)  # load in float16