update qwen model

Signed-off-by: tk <[email protected]>
ssbuild · Dec 2, 2023 · 524f307 · 524f307
1 parent 4fc0c8e
commit 524f307
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/deep_training/nlp/models/qwen/modeling_qwen.py b/src/deep_training/nlp/models/qwen/modeling_qwen.py
@@ -537,7 +537,8 @@ def forward(
                         -1, -1, causal_mask.size(2), -1
                     )
                     if causal_mask is not None:
-                        attention_mask.masked_fill(~causal_mask, torch.finfo(query.dtype).min)
+                        # attention_mask.masked_fill(~causal_mask, torch.finfo(query.dtype).min)
+                        attention_mask.masked_fill(~causal_mask, torch.finfo(attention_mask.dtype).min)
                 else:
                     attention_mask = causal_mask
                 attn_output = F.scaled_dot_product_attention(