Update llama_diffusion_model.py

llama_diffusion_model.py

@@ -122,12 +122,12 @@ class CustomTransformerModel(PreTrainedModel):
         device = input_ids.device
         
         masking_type = getattr(self.config, "masking_type", "bidirectional_masked")
-        if self.config.masking_type == 'bidirectional':
+        if masking_type == 'bidirectional':
             base_mask = torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)
-        elif self.config.masking_type == 'bidirectional_masked':
+        elif masking_type == 'bidirectional_masked':
             base_mask = torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)
             base_mask.fill_diagonal_(False)
-        elif self.config.masking_type == 'unidirectional':
+        elif masking_type == 'unidirectional':
             base_mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))
         else:
             raise ValueError(f"Unknown masking type: {self.config.masking_type}")