Update modeling_diffusion.py

modeling_diffusion.py

@@ -5,6 +5,15 @@ class DiffusionTextModel(nn.Module, PyTorchModelHubMixin):
     def __init__(self, vocab_size, max_seq_len, max_time_steps,
                  embed_dim=128, n_layers=4, n_heads=4):
         super().__init__()
+        self.config = {
+            "vocab_size": vocab_size,
+            "max_seq_len": max_seq_len,
+            "max_time_steps": max_time_steps,
+            "embed_dim": embed_dim,
+            "n_layers": n_layers,
+            "n_heads": n_heads
+        }
+
         self.token_emb = nn.Embedding(vocab_size, embed_dim)
         self.pos_emb   = nn.Embedding(max_seq_len, embed_dim)
         self.time_emb  = nn.Embedding(max_time_steps+1, embed_dim)
@@ -23,4 +32,4 @@ class DiffusionTextModel(nn.Module, PyTorchModelHubMixin):
         tim = self.time_emb(t).unsqueeze(1).expand(B, L, -1)
         h = tok + pos + tim
         h = self.transformer(h.transpose(0,1)).transpose(0,1)
-        return self.out(h)
+        return self.out(h)