Update modeling_gpt_refact.py

modeling_gpt_refact.py

@@ -151,7 +151,7 @@ class Attention(nn.Module):
         upcast = dtype != softmax_dtype
         unscale = self.layer_idx + 1 if self.scale_attention_softmax_in_fp32 and upcast else 1
 
-        attn_weights = alibi + torch.matmul(query * self.scale, key)
+        attn_weights = (alibi + torch.matmul(query * self.scale, key)).to(query.dtype)
 
         if upcast:
             if attention_mask is None: