Fix small typo in JambaDecoder (#24)

- Fix small typo in JambaDecoder (3e2e3bf1a6540e78b3d6e8f58adbbab384d6d32d)


Co-authored-by: Moshe Berchansky <[email protected]>

modeling_jamba.py

@@ -1053,7 +1053,7 @@ class JambaMambaMixer(nn.Module):
             ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :]      # [batch, intermediade_size, ssm_state]
             scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1))  # [batch, intermediade_size, 1]
             scan_outputs.append(scan_output[:, :, 0])
-        scan_output = torch.stack(scan_outputs, dim=-1)                                # [batch, seq_len, intermediade_size]
+        scan_output = torch.stack(scan_outputs, dim=-1)                                # [batch, intermediade_size, seq_len]
         scan_output = scan_output + (hidden_states * self.D[None, :, None])
         scan_output = (scan_output * self.act(gate))