Relax assert to allow both bf16 and fp16

2022-09-11 12:09:43 -07:00 · 2022-09-11 12:09:43 -07:00 · 13403e8115
commit 13403e8115
parent 64f42cd057
1 changed files with 1 additions and 1 deletions
--- a/flash_attn/flash_attention.py
+++ b/flash_attn/flash_attention.py
@ -34,7 +34,7 @@ class FlashAttention(nn.Module):
            key_padding_mask: a bool tensor of shape (B, S)
        """
        assert not need_weights
-        assert qkv.dtype == torch.float16
+        assert qkv.dtype in [torch.float16, torch.bfloat16]
        assert qkv.is_cuda

        if cu_seqlens is None: