Make attention mask float

llama_diffusion_model.py

@@ -133,6 +133,8 @@ class CustomTransformerModel(PreTrainedModel):
             raise ValueError(f"Unknown masking type: {self.config.masking_type}")
 
         attention_mask = base_mask.unsqueeze(0).unsqueeze(1).expand(batch_size, 1, seq_len, seq_len).clone()
+        attention_mask = attention_mask.to(dtype=torch.float32)  # required for SDPA and Flash attention
+
 
         with autocast("cuda", dtype=torch.float16):
             outputs = self.llama(