squall/flash-attention
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix typo in RotaryEmbedding forward output type (#666)

Browse Source
This commit is contained in:
Antony Frolov 2023-11-09 22:43:02 +03:00 committed by GitHub
parent 83aef842be
commit 3566596ad8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
flash_attn/layers/rotary.py
View File

@ -417,7 +417,7 @@ class RotaryEmbedding(torch.nn.Module):
kv: Optional[torch.Tensor] = None,
seqlen_offset: Union[int, torch.Tensor] = 0,
max_seqlen: Optional[int] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
"""
qkv: (batch, seqlen, 3, nheads, headdim) if kv is none,
else it's just q of shape (batch, seqlen, nheads, headdim)