suno-ai/bark
1
0
Fork 0
mirror of https://github.com/suno-ai/bark.git synced 2026-10-11 22:09:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #364 from no2chem/flash_attention

Browse Source 
fix: perf: use sdpa if dropout > 0 on fine model
This commit is contained in:
Martin Camacho
2023-07-05 08:58:08 -04:00
committed by GitHub
parent 7280e4e31c b56c8df48a
commit 599fed040e
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
bark/model_fine.py
+2 -2
View File
@@ -26,9 +26,9 @@ class NonCausalSelfAttention(nn.Module):
self.n_head = config.n_head
self.n_embd = config.n_embd
self.dropout = config.dropout
# flash attention make GPU go brrrrr but support is only in PyTorch nightly and still a bit scary
# flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
self.flash = (
hasattr(torch.nn.functional, "scaled_dot_product_attention") and self.dropout == 0.0
hasattr(torch.nn.functional, "scaled_dot_product_attention")
)
def forward(self, x):