Fix CPU Fallback for NewAttention with xformers BlockDiagonalMask

modeling.py

@@ -445,6 +445,10 @@ class NewAttention(nn.Module):
 
         if use_memory_efficient_attention is None:
             use_memory_efficient_attention = self.config.use_memory_efficient_attention
+
+        if not torch.cuda.is_available() or (hasattr(config, 'device') and config.device == 'cpu'):
+            use_memory_efficient_attention = False
+            
         self.use_memory_efficient_attention = use_memory_efficient_attention
         self.memory_efficient_attention = None if xops is None else xops.memory_efficient_attention
         if self.use_memory_efficient_attention:
@@ -489,6 +493,9 @@ class NewAttention(nn.Module):
             key_states = pad_input(key_states.squeeze(), *padding_inputs)
             value_states = pad_input(value_states.squeeze(), *padding_inputs)
 
+        if self.use_memory_efficient_attention and not hidden_states.is_cuda:
+            self.use_memory_efficient_attention = False
+        
         if self.use_memory_efficient_attention:
             assert self.memory_efficient_attention is not None, "xformers is not loaded"
             assert output_attentions is False, "memory_efficient_attention do not output attentions"
@@ -534,6 +541,12 @@ class NewAttention(nn.Module):
 
         attention_scores = attention_scores / math.sqrt(self.attention_head_size)
         if attention_bias is not None:
+           if hasattr(attention_bias, 'materialize'):
+                # If it's a BlockDiagonalMask, materialize it to a tensor
+                attention_bias = attention_bias.materialize(
+                    (attention_scores.shape[0], attention_scores.shape[1], 
+                    attention_scores.shape[2], attention_scores.shape[3])
+                )
             # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
             attention_scores = attention_scores + attention_bias