Update modeling_motif.py

modeling_motif.py

@@ -1032,9 +1032,9 @@ class MotifFlashAttention2(MotifAttention):
                                         causal=causal)
             return attn_out
         else:
-            attn_out = _flash_attention_forward(query_states,
-                                            key_states,
-                                            value_states,
+            attn_out = _flash_attention_forward(query_states.bfloat16(),
+                                            key_states.bfloat16(),
+                                            value_states.bfloat16(),
                                             attention_mask,
                                             q_len,
                                             position_ids=position_ids,
@@ -1044,7 +1044,7 @@ class MotifFlashAttention2(MotifAttention):
                                             softmax_scale=scale_factor,
                                             use_top_left_mask=self._flash_attn_uses_top_left_mask)
             #logger.info(attn_out)
-            return attn_out
+            return attn_out.float()
 
     def forward(
             self,