Update modeling_diffusion.py

modeling_diffusion.py

@@ -4,7 +4,7 @@ from transformers import PreTrainedModel
 from .configuration_diffusion import DiffusionConfig
 
 class DiffusionTransformer(nn.Module):
-    def __init__(self, vocab_size, d_model=256, n_heads=8, n_layers=6, dropout=0.1, max_len=140, num_steps=1000):
+    def __init__(self, vocab_size, d_model=512, n_heads=16, n_layers=6, dropout=0.1, max_len=140, num_steps=1000):
         super().__init__()
         self.vocab_size = vocab_size
         self.max_len = max_len