How do I input a four-dimensional tensor into Mamba2?

[NicoleDyson]

For example, I have a four-dimensional tensor with shape [2, 56, 56, 96], where the four dimensions correspond to batch_size, height, width, and channels. If I directly set the parameters as follows:

d_model = width * height  # This is incorrect
d_state = 64,
d_conv = `4,`
expand = 2,

and create the Mamba2 module, it results in an error:

File "/root/anaconda3/envs/test/lib/python3.10/site-packages/mamba_ssm/modules/mamba2.py", line 157, in forward
batch, seqlen, dim = u.shape
ValueError: too many values to unpack (expected 3)

[AlwaysFHao]

Mamba's input requires the format of [batch, seq_len, dim]. You can refer to the method of Vision Transformer to transform the data of [batch, c, h, w] into [batch, dim, patch, patch] using Conv2d, and then transpose it to the shape of [batch, patch * patch, dim] for input. The above satisfies seq_len=(h / patch) ^ 2.