Update model.py

model.py

@@ -30,6 +30,9 @@ class minGRU(nn.Module):
     def __init__(self, d_model, d_inner):
         super().__init__()
 
+        self.d_model = d_model
+        self.d_inner = d_inner
+
         self.hidden_proj = nn.Linear(d_model, d_inner, bias=False)
         self.gate_proj   = nn.Linear(d_model, d_inner, bias=False)
         self.out_proj    = nn.Linear(d_inner, d_model, bias=False)
@@ -100,6 +103,12 @@ class RMSNorm(nn.Module):
 class minGRULM(nn.Module):
     def __init__(self, vocab_size, d_model, d_inner, n_layers):
         super().__init__()
+
+        self.vocab_size = vocab_size
+        self.d_model    = d_model
+        self.d_inner    = d_inner
+        self.n_layers   = n_layers
+        
         self.embed = nn.Embedding(vocab_size, d_model)
 
         self.layers = nn.ModuleList([])