Update modeling_quiet.py

modeling_quiet.py

@@ -449,9 +449,10 @@ class QuietFlashAttention2(QuietAttention):
             key_states = key_states.to(target_dtype)
             value_states = value_states.to(target_dtype)
         # 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)
+        query_states = query_states.reshape(bsz, -1, self.num_heads, self.head_dim)
+        key_states = key_states.reshape(bsz, -1, self.num_key_value_heads, self.head_dim)
+        value_states = value_states.reshape(bsz, -1, self.num_key_value_heads, self.head_dim)
+
 
         attn_output = self._flash_attention_forward(
             query_states,