Merge branch 'state-spaces:main' into feat/add-cu_seqlens

Author: zigzagcai

mamba_ssm/modules/mamba2_simple.py

@@ -167,6 +167,7 @@ def forward(self, u, seq_idx=None):
                 xBC = self.act(
                     self.conv1d(xBC.transpose(1, 2)).transpose(1, 2)
                 )  # (B, L, self.d_inner + 2 * ngroups * d_state)
+                xBC = xBC[:, :seqlen, :]
             else:
                 xBC = causal_conv1d_fn(
                     x=xBC.transpose(1, 2),