docs(transformer): 更新process_attention_mask函数文档

2025-11-05 23:41:11 +08:00 · 2025-11-05 23:41:11 +08:00 · 805773c7fe
parent 7ccc4ab9ac
commit 805773c7fe
1 changed files with 1 additions and 1 deletions
--- a/khaosz/model/transformer.py
+++ b/khaosz/model/transformer.py
@ -18,10 +18,10 @@ def process_attention_mask(
    Create attention mask for GQA
    Args:
        seq_mask (Tensor): A tensor indicating whether each position is valid or not.
+        input_tensor (Tensor): The input tensor.
        start_pos (int): The starting position of the sequence.
        seq_len (int): The length of the sequence.
        is_causal (bool): Whether the attention is causal or not.
-        device (torch.device): The device to use.
    Returns:
        Tensor: The attention mask tensor.
    """