From b56c8df48a23cbf9726b36f21f9a586a571c4241 Mon Sep 17 00:00:00 2001 From: Michael Wei Date: Tue, 20 Jun 2023 11:37:21 -0700 Subject: [PATCH] fix: use sdpa if dropout > 0 on fine model --- bark/model_fine.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bark/model_fine.py b/bark/model_fine.py index 6179a85..e51d3e6 100644 --- a/bark/model_fine.py +++ b/bark/model_fine.py @@ -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):