diff --git a/multitalk/multitalk.py b/multitalk/multitalk.py
index 5922ac5..c87f906 100644
--- a/multitalk/multitalk.py
+++ b/multitalk/multitalk.py
@@ -256,13 +256,11 @@ class SingleStreamAttention(nn.Module):
         N_t, N_h, N_w = shape
 
         x_extra = None
-        try:
-            x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
-        except:
+        if x.shape[0] * N_t != 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)
+        x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t)
 
         # get q for hidden_state
         B, N, C = x.shape
@@ -369,13 +367,13 @@ class SingleStreamMultiAttention(SingleStreamAttention):
             return super().forward(x, encoder_hidden_states, shape)
 
         N_t, N_h, N_w = 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]:
+        if x.shape[0] * N_t != 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)
 
         # Query projection
         B, N, C = x.shape