Changed back to bidirectional attention

llama_diffusion_model.py

@@ -192,7 +192,7 @@ class CustomTransformerModel(PreTrainedModel):
         self.llama.resize_token_embeddings(config.vocab_size)
 
         for i, layer in enumerate(self.llama.model.layers):
-            layer.self_attn = BidirectionalLlamaAttention(layer.self_attn, masking='bidirectional_masked')
+            layer.self_attn = BidirectionalLlamaAttention(layer.self_attn, masking='bidirectional')
 
         # Freeze Llama to retain pre-trained knowledge
         for param in self.llama.parameters():