update attn_weights add mask dtype not match error

2025-10-11 11:13:10 +08:00
parent dfeabc33f8
commit dcdafaf8ac
1 changed files with 1 additions and 1 deletions
@@ -225,7 +225,7 @@ impl AttentionNobias {
                let attn_weights = (attn_weights * scale)?;
                let attn_weights = match attention_mask {
                    None => attn_weights,
-                    Some(mask) => attn_weights.broadcast_add(mask)?,
+                    Some(mask) => attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)?,
                };
                let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
                let attn_weights = attn_weights.matmul(&value_states)?;