Torchsummary not working

[avocardio]

When trying to inspect a model via torchsummary, it returns this error.

Traceback (most recent call last):
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torchsummary/torchsummary.py", line 140, in summary
    _ = model.to(device)(*x, *args, **kwargs)  # type: ignore[misc]
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/maximilian/Documents/NR/NR-SSL/../NR-SSL/utils/pt_utils.py", line 56, in wrapper
    result = func(self, *args, **kwargs)
  File "/Users/maximilian/Documents/NR/NR-SSL/models/ecog/S_XTransformer.py", line 67, in forward
    x = self.encoder(x)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torch/nn/modules/container.py", line 217, in forward
    input = module(input)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/x_transformers/x_transformers.py", line 1408, in forward
    x = self.attn_layers(x, mask = mask, mems = mems, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/x_transformers/x_transformers.py", line 1153, in forward
    out, inter = block(x, mask = mask, context_mask = self_attn_context_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, rotary_pos_emb = rotary_pos_emb, prev_attn = prev_attn, mem = layer_mem)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/x_transformers/x_transformers.py", line 835, in forward
    attn_bias = rel_pos(i, j)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1547, in _call_impl
    hook_result = hook(self, args, result)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torchsummary/torchsummary.py", line 304, in hook
    info.input_size = info.calculate_size(inputs, batch_dim)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torchsummary/layer_info.py", line 82, in calculate_size
    size = nested_list_size(inputs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torchsummary/layer_info.py", line 58, in nested_list_size
    return nested_list_size(inputs[0])
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torchsummary/layer_info.py", line 55, in nested_list_size
    if hasattr(inputs[0], "size") and callable(inputs[0].size):
TypeError: 'int' object is not subscriptable

The above exception was the direct cause of the following exception:

...
  File "/opt/homebrew/Caskroom/miniforge/base/envs/light/lib/python3.8/site-packages/torchsummary/torchsummary.py", line 143, in summary
    raise RuntimeError(
RuntimeError: Failed to run torchsummary. See above stack traces for more details. Executed layers up to: [AbsolutePositionalEmbedding: 3-1, Embedding: 4-1, TokenEmbedding: 3-2, Embedding: 4-2, Identity: 3-3, Dropout: 3-4, Identity: 3-5]