Update modeling_motif.py

modeling_motif.py

@@ -493,7 +493,7 @@ class MotifFlashAttention2(MotifAttention):
     def _compute_attention(self, query_states, key_states, value_states, attention_mask, q_len, position_ids,
                            dropout_rate, sliding_window):
         """Flash Attention 2 implements"""
-        
+        _input_type = query_states.dtype        
         scale_factor = 1.0 / math.sqrt(self.head_dim)
         if not self._flash_attn_uses_top_left_mask:
             causal = self.is_causal
@@ -511,7 +511,7 @@ class MotifFlashAttention2(MotifAttention):
                                             is_causal=True,
                                             softmax_scale=scale_factor,
                                             use_top_left_mask=self._flash_attn_uses_top_left_mask)
-        return attn_out.float()
+        return attn_out.to(_input_type)
 
     def forward(
             self,