Torchinfo issue

So, I’m working on my research and I have build a model and I’m facing a problem while printing model’s summary.

The model takes a (3,224,224) shape image with batch size = 2. When I pass a batch through the model, it passes without any problem. I wanted to see the model summary so, I used torchinfo's summary and I face this error:

TypeError: Model contains a layer with an unsupported input or output type: 0, type: <class 'int'>

I can show the last traceback call which led to this error:


D:\Anaconda\envs\ffconvolution\lib\site-packages\torchinfo\torchinfo.py in hook(***failed resolving arguments***)
    507         info.input_size = info.calculate_size(inputs, batch_dim)
--> 508         info.output_size = info.calculate_size(outputs, batch_dim)
    509         info.executed = True

D:\Anaconda\envs\ffconvolution\lib\site-packages\torchinfo\layer_info.py in calculate_size(inputs, batch_dim)
    101         else:
--> 102             raise TypeError(
    103                 "Model contains a layer with an unsupported input or output type: "

I’m not sure what’s going wrong. Give suggestions.