Error: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!"

Question

I'm running this code on Azure machine learning notebook.

import os
import torch
import gradio as gr
from vilmedic import AutoModel
from vilmedic.blocks.scorers import RadGraph
import glob

model, processor = AutoModel.from_pretrained("rrg/baseline-mimic")
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = model.to(device)
radgraph = RadGraph(cuda=-1)

.....

def run(image, beam_size, num_return_sequences, include_words, exclude_words, do_radgraph):
    if image is None:
        return {}, '<b>Please select an image</b>'  # , ""
    if num_return_sequences > beam_size:
        return {}, '<b>"Beam size"</b> must be greater or equal than <b>"Number of generated reports"</b>'  # , ""

    try:
        include_words_ids, include_words = get_token_from_strings(include_words)
        exclude_words_ids, exclude_words = get_token_from_strings(exclude_words)

        if include_words_ids is not None and [3] in include_words_ids:
            return {}, '<b>"' + include_words[
                include_words_ids.index([3])] + '"</b> is not in the vocabulary"</b>'  # , ""

        with torch.no_grad():
            batch = processor.inference(image=[
                [image]
            ])
            batch_size = 1
            encoder_output, encoder_attention_mask = model.encode(**batch)
            expanded_idx = torch.arange(batch_size).view(-1, 1).repeat(1, beam_size).view(-1)
            input_ids = torch.ones((len(batch["images"]), 1), dtype=torch.long)
            if torch.cuda.is_available():
                expanded_idx = expanded_idx.cuda()
                print("89 line")
                input_ids = input_ids.cuda()

            # Using huggingface generate method
            hyps = model.dec.generate(
                input_ids=input_ids * model.dec.config.bos_token_id,
                encoder_hidden_states=encoder_output.index_select(0, expanded_idx),
                encoder_attention_mask=encoder_attention_mask.index_select(0, expanded_idx),
                num_return_sequences=num_return_sequences,
                max_length=processor.tokenizer_max_len,
                num_beams=beam_size,
                bad_words_ids=exclude_words_ids,
                force_words_ids=include_words_ids,
            )

            # Decode
            hyps = [processor.tokenizer.decode(h, skip_special_tokens=True, clean_up_tokenization_spaces=False) for h in
                    hyps]

            # RadGraph
            if do_radgraph:
                radgraph_annots = [radgraph(hyps=[h], refs=[h])[-1][0]["entities"] for h in hyps]
                # Find entites : Radgraph
                new_hyp_strs = []
                for hyp_str, radgraph_annot in zip(hyps, radgraph_annots):
                    values = radgraph_annot.values()
                    new_hyp_str = hyp_str.split()
                    for v in values:
                        new_hyp_str[v["start_ix"]] = highlight_radgraph_entities(v["tokens"], v["label"])
                    new_hyp_strs.append(' '.join(new_hyp_str))
            else:
                new_hyp_strs = hyps

            # Find user entites
            if include_words is not None:
                for w in include_words:
                    new_hyp_strs = [h.replace(w, highlight_word(w, "user")) for h in new_hyp_strs]

            # Formating
            new_hyp_strs = ["<p><b>Hypothesis {}:</b> <br/> {} </p>" \
                            "".format(i + 1, h) for i, h in enumerate(new_hyp_strs)] + (
                               ["<br/><br/><i>Anat: anatomy<br/>"
                                "OBS: observation<br/>"
                                "DA: definitely absent<br/>"
                                "DP: definitely present</i>"] if do_radgraph else [""])

            # Params
            out_json = {
                "beam size": beam_size, "number of generated reports": num_return_sequences,
                "included words": include_words, "excluded words": exclude_words, "show radgraph": do_radgraph
            }

            return out_json, str(''.join(new_hyp_strs))  # , str(refs[os.path.basename(image)])

    except Exception as e:
        print(e)
        return {}, "<b>An error occured, try again..."

The full code is here:

https://huggingface.co/spaces/StanfordAIMI/radiology_report_generation/blob/main/app.py

It keeps giving me this error when I upload an image and click the submit button: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!. I'm using an azure computation with following specifications: Virtual machine size: Standard_NV6 (6 cores, 56 GB RAM, 380 GB disk) Processing unit: GPU - 1 x NVIDIA Tesla M60