Update chatNT.py

chatNT.py

@@ -1324,6 +1324,9 @@ class MultiHeadAttention(nn.Module):
         attention_weights = attention_weights / sqrt_key_size
         if attention_mask is not None:
             attention_weights = torch.where(attention_mask, attention_weights, -1e30)
+
+        attention_weights = attention_weights.to(value_heads.dtype)
+
         if attention_weight_bias is not None:
             attention_weights = F.softmax(
                 attention_weights + attention_weight_bias, dim=-1