mirror of
https://github.com/huggingface/candle.git
synced 2025-06-16 02:38:10 +00:00
Softmax numerical stability. (#267)
* Softmax numerical stability. * Fix the flash-attn test.
This commit is contained in:
@ -333,7 +333,7 @@ impl BertSelfAttention {
|
||||
let attention_scores = (attention_scores / (self.attention_head_size as f64).sqrt())?;
|
||||
let attention_probs = {
|
||||
let _enter_sm = self.span_softmax.enter();
|
||||
attention_scores.softmax(candle::D::Minus1)?
|
||||
candle_nn::ops::softmax(&attention_scores, candle::D::Minus1)?
|
||||
};
|
||||
let attention_probs = self.dropout.forward(&attention_probs)?;
|
||||
|
||||
|
Reference in New Issue
Block a user