Update model.py

src/models/model.py

@@ -252,24 +252,34 @@ class LightningModel(LightningModule):
         no_decay = ["bias", "LayerNorm.weight"]
         optimizer_grouped_parameters = [
             {
-                "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+                "params": [
+                    p
+                    for n, p in model.named_parameters()
+                    if not any(nd in n for nd in no_decay)
+                ],
                 "weight_decay": self.weight_decay,
             },
             {
-                "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
+                "params": [
+                    p
+                    for n, p in model.named_parameters()
+                    if any(nd in n for nd in no_decay)
+                ],
                 "weight_decay": 0.0,
             },
         ]
-        optimizer = AdamW(optimizer_grouped_parameters, lr=self.learning_rate, eps=self.adam_epsilon)
+        optimizer = AdamW(
+            optimizer_grouped_parameters, lr=self.learning_rate, eps=self.adam_epsilon
+        )
         self.opt = optimizer
         return [optimizer]
 
 
 class Summarization:
-    """ Custom Summarization class """
+    """Custom Summarization class"""
 
     def __init__(self) -> None:
-        """ initiates Summarization class """
+        """initiates Summarization class"""
         pass
 
     def from_pretrained(self, model_type="t5", model_name="t5-base") -> None: