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/