fix(transformer): 调整注意力掩码处理逻辑

2025-09-29 11:31:42 +08:00 · 2025-09-29 11:31:42 +08:00 · 8206c7855e
parent 816bc78894
commit 8206c7855e
1 changed files with 5 additions and 5 deletions
--- a/khaosz/core/transformer.py
+++ b/khaosz/core/transformer.py
@ -91,11 +91,11 @@ def process_attention_mask(
        Tensor: The attention mask tensor.
    """
    
-    if start_pos != 0 and seq_mask is None:
+    if seq_mask is None:
+        if start_pos != 0:
            # for single prompt chat
            seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
-    
-    if seq_mask is None:
+        else:
            return None
    
    if seq_mask.dim() > 2: