torchsummary error

[halo34]

Hi, I have some trouble with monai and torchsummary.
unet_attention = monai.networks.nets.AttentionUnet( spatial_dims=2, in_channels=3, out_channels=2, channels=(16, 32, 64, 128, 256), strides=(2, 2, 2, 2) ) unetr = monai.networks.nets.UNETR( in_channels=3, out_channels=2, img_size=image_size, spatial_dims=2, ) torchsummary.summary(unet_attention, input_size=(3,256,256)) torchsummary.summary(unetr, input_size=(3,256,256))
I got the following error:
Traceback (most recent call last): File "create_summary.py", line 33, in <module> torchsummary.summary(unetr, input_size=(3.,256,256)) File "/opt/homebrew/lib/python3.11/site-packages/torchsummary/torchsummary.py", line 72, in summary model(*x) File "/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl return forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/lib/python3.11/site-packages/monai/networks/nets/unetr.py", line 200, in forward x, hidden_states_out = self.vit(x_in) ^^^^^^^^^^^^^^ File "/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1547, in _call_impl hook_result = hook(self, args, result) ^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/lib/python3.11/site-packages/torchsummary/torchsummary.py", line 22, in hook summary[m_key]["output_shape"] = [ ^ File "/opt/homebrew/lib/python3.11/site-packages/torchsummary/torchsummary.py", line 23, in <listcomp> [-1] + list(o.size())[1:] for o in output ^^^^^^ AttributeError: 'list' object has no attribute 'size'

[KumoLiu]

Hi @halo34, I think it's a limitation in torchsummary, could you please try torch-summary?
pip install torch-summary==1.4.4

Hope it helps, thanks!

[halo34]

Thank you. It works with torchinfo too.