Swinv2 visualization fails with TypeError: '>' not supported between instances of 'int' and 'torch.Size'

[f-fuchs]

Hey,

I am trying to visualize the swinv2 model (combined with FeatureListNet) of the timm libary but sadly I get the following error.

), shape=(torch.Size([32, 8, 32, 256]), torch.Size([32, 4, 16, 512]), torch.Size([32, 2, 8, 1024]), 1), scale_xy=1, x=10, max_xy=2000
Traceback (most recent call last):
  File "/home/fuchsfa/feature-exploration/model_vis.py", line 56, in <module>
    img = visualtorch.layered_view(model, input_shape=input_shape, legend=True)
  File "/home/fuchsfa/feature-exploration/.venv/lib/python3.13/site-packages/visualtorch/layered.py", line 124, in layered_view
    layer_y, layer_types, boxes, x_off, img_height, max_right = _create_architecture(
                                                                ~~~~~~~~~~~~~~~~~~~~^
        layers,
        ^^^^^^^
    ...<17 lines>...
        padding,
        ^^^^^^^^
    )
    ^
  File "/home/fuchsfa/feature-exploration/.venv/lib/python3.13/site-packages/visualtorch/layered.py", line 323, in _create_architecture
    x = min(max(shape[1] * scale_xy, x), max_xy)
            ~~~^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: '>' not supported between instances of 'int' and 'torch.Size'

It seems to not be iterating over the layers correctly becaues it thinks this is a layer

key='FeatureListNet-498', layer=FeatureListNet(                                                                                                                                                                                      
  (patch_embed): PatchEmbed(                                                                                                                                                                                                         
    (proj): Conv2d(1, 128, kernel_size=(4, 4), stride=(4, 4))                                                                                                                                                                        
    (norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)                                                                                                                                                                    
  )                                                                                                                                                                                                                                  
  (layers_0): SwinTransformerV2Stage(                                                                                                                                                                                                
    (downsample): Identity()       
  ....