Fix bug in mamba2_simple slow path conv1d state-spaces#393

mzusman · Jul 22, 2024 · a71bb5a · a71bb5a
1 parent c0a00bd
commit a71bb5a
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/mamba_ssm/modules/mamba2_simple.py b/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),