[bugfix] Initialize attention bias on the same device as Query/Key/Value

modeling.py

@@ -910,7 +910,7 @@ class NewModel(NewPreTrainedModel):
 
         batch_size, seq_length = input_shape
         if unpad_inputs and self.config.use_memory_efficient_attention:
-            attention_bias = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(length)
+            attention_bias = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(length, device=embedding_output.device)
         else:
             # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
             # ourselves in which case we just need to make it broadcastable to all heads.