RuntimeError when running on CPU

[jose-alanaai]

When running the examples from the colab, they work fine on GPU but I get the following runtime error when trying to run on CPU:

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.HalfTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

Below is a snippet of code to reproduce the error:

import torch
import requests

from PIL import Image
from lavis.models import load_model_and_preprocess

device = torch.device('cpu')
with torch.inference_mode():
    model, vis_processors, _ = load_model_and_preprocess(
        name='blip2_t5',
        model_type='pretrain_flant5xl',
        is_eval=True,
        device=device
    )

img_url = 'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
image = vis_processors["eval"](raw_image).unsqueeze(0).to(device)

model.generate({"image": image, "prompt": "Question: which city is this? Answer:"})

I'm wondering whether this is due to some version mismatch between any of the libraries or whether perhaps there is somewhere in the code where the device is being ignored? The versions of torch and transformers are as follows:

torch                    1.12.1+cu113
torchvision              0.13.1+cu113
transformers             4.26.1

P.S.: This issue is similar to https://github.com/salesforce/LAVIS/issues/145 but the opposite behaviour...