Update modeling_quiet.py

modeling_quiet.py

@@ -448,11 +448,10 @@ class QuietFlashAttention2(QuietAttention):
             query_states = query_states.to(target_dtype)
             key_states = key_states.to(target_dtype)
             value_states = value_states.to(target_dtype)
-
-        # Reashape to the expected shape for Flash Attention
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
+        # Reshape to the expected shape for Flash Attention
+        query_states = query_states.view(bsz, -1, self.num_heads, self.head_dim)
+        key_states = key_states.view(bsz, -1, self.num_key_value_heads, self.head_dim)
+        value_states = value_states.view(bsz, -1, self.num_key_value_heads, self.head_dim)
 
         attn_output = self._flash_attention_forward(
             query_states,
@@ -462,7 +461,7 @@ class QuietFlashAttention2(QuietAttention):
             q_len,
             dropout=dropout_rate,
             use_sliding_windows=use_sliding_windows,
-        )
+)
 
         attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
         attn_output = self.o_proj(attn_output)