Error with models with non tensor inputs

[Krish-Sahu]

AttributeError Traceback (most recent call last)
in <cell line: 2>()
1 dummy_inputs = torch.rand(1, 3, 224, 224)
----> 2 flops, params = flopth(matcher, inputs=({"image0": feats0, "image1": feats1}))
3 print(flops, params)

1 frames
/usr/local/lib/python3.10/dist-packages/flopth/init.py in flopth(model, in_size, inputs, dtype, param_dict, show_detail, bare_number)
212 input_list.append(x)
213
--> 214 mv = ModelViewer(model, input_list, param_dict, dtype)
215
216 if show_detail:

/usr/local/lib/python3.10/dist-packages/flopth/model_viewer.py in init(self, model, input_list, param_dict, dtype)
20 if torch.cuda.is_available():
21 for i in range(len(input_list)):
---> 22 input_list[i] = input_list[i].cuda()
23
24 self._model = self._model.cuda()

AttributeError: 'str' object has no attribute 'cuda'

The reason for this is in the ModelViewer class in the model_viewer.py file -

    if torch.cuda.is_available():
        for i in range(len(input_list)):
            input_list[i] = input_list[i].cuda()

        self._model = self._model.cuda()
    else:
        device = next(self._model.parameters()).device
        for i in range(len(input_list)):
            input_list[i] = input_list[i].to(device)

Here if the input is a dict it will raise error as it has no .cuda() attribute/