diff --git a/multitalk/multitalk.py b/multitalk/multitalk.py
index 610b829..581b43f 100644
--- a/multitalk/multitalk.py
+++ b/multitalk/multitalk.py
@@ -253,8 +253,14 @@ class SingleStreamAttention(nn.Module):
         self.attention_mode = attention_mode
 
     def forward(self, x: torch.Tensor, encoder_hidden_states: torch.Tensor, shape=None, enable_sp=False, kv_seq=None) -> torch.Tensor:
-       
         N_t, N_h, N_w = shape
+
+        x_extra = None
+        if x.shape[0] != encoder_hidden_states.shape[0]:
+            x_extra = x[:, -N_h * N_w:, :]
+            x = x[:, :-N_h * N_w, :]
+            N_t = N_t - 1
+      
         x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
 
         # get q for hidden_state
@@ -282,17 +288,18 @@ class SingleStreamAttention(nn.Module):
             encoder_v.transpose(1, 2),
             attention_mode=self.attention_mode
             )
-        #x = torch.nn.functional.scaled_dot_product_attention(
-        #    q, encoder_k, encoder_v, attn_mask=None, is_causal=False, dropout_p=0.0)
 
         # linear transform
         x_output_shape = (B, N, C)
         #x = x.transpose(1, 2) 
         x = x.reshape(x_output_shape) 
         x = self.proj(x)
-        x = self.proj_drop(x)
+        x = self.proj_drop(x)        
 
         x = rearrange(x, "(B N_t) S C -> B (N_t S) C", N_t=N_t)
+    
+        if x_extra is not None:
+            x = torch.cat([x, torch.zeros_like(x_extra)], dim=1)
 
         return x
     
@@ -363,6 +370,12 @@ class SingleStreamMultiAttention(SingleStreamAttention):
         N_t, _, _ = shape
         x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
 
+        x_extra = None
+        if x.shape[0] != encoder_hidden_states.shape[0]:
+            x_extra = x[:, -N_h * N_w:, :]
+            x = x[:, :-N_h * N_w, :]
+            N_t = N_t - 1
+
         # Query projection
         B, N, C = x.shape
         q = self.q_linear(x)
@@ -473,5 +486,7 @@ class SingleStreamMultiAttention(SingleStreamAttention):
 
         # Restore original layout
         x = rearrange(x, "(B N_t) S C -> B (N_t S) C", N_t=N_t)
+        if x_extra is not None:
+            x = torch.cat([x, torch.zeros_like(x_extra)], dim=1)
 
         return x
\ No newline at end of file