Update model.py

src/models/model.py

@@ -165,6 +165,7 @@ class LightningModel(LightningModule):
         self.tokenizer = tokenizer
         self.output = output
         self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
 
     def forward(self, input_ids, attention_mask, decoder_attention_mask, labels=None):
         """ forward step """