Add reshape apply to input projection

jloveric · Dec 30, 2023 · 62330c6 · 62330c6
1 parent 80161b1
commit 62330c6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/language_interpolation/state_space_network.py b/language_interpolation/state_space_network.py
@@ -181,7 +181,7 @@ class Mamba, https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/m
 
         (b, l, d) = x.shape
 
-        x_and_res = self.in_proj(x)  # shape (b, l, 2 * d_in)
+        x_and_res = reshape_apply(x, self.in_proj)  # shape (b, l, 2 * d_in)
         (x, res) = x_and_res.split(split_size=[self.args.d_inner, self.args.d_inner], dim=-1)
 
         x = rearrange(x, 'b l d_in -> b d_in l')