[fix] fix issues/382 修复scaled_dot_product_attention方法不支持cpu、mps的问题

2025-04-29 17:36:30 +08:00 · 2025-04-29 17:36:30 +08:00 · 375c22ec42
parent 4a1a9bec8a
commit 375c22ec42
1 changed files with 1 additions and 0 deletions
--- a/model/model_minimind.py
+++ b/model/model_minimind.py
@ -141,6 +141,7 @@ class Attention(nn.Module):
        self.resid_dropout = nn.Dropout(args.dropout)
        self.dropout = args.dropout
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn and torch.cuda.is_available()
        # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")

    def forward(self,