Merge pull request suno-ai#364 from no2chem/flash_attention

fix: perf: use sdpa if dropout > 0 on fine model

Author: mcamac

bark/model_fine.py

@@ -26,9 +26,9 @@ def __init__(self, config):
         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):