Pre-commit fixes, optimizer loading fix

2 files changed, 5 insertions, 7 deletions

diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py
index 1004f48..11d1eb1 100644
--- a/text_recognizer/models/base.py
+++ b/text_recognizer/models/base.py
@@ -15,7 +15,7 @@ class LitBaseModel(pl.LightningModule):
 
     def __init__(
         self,
-        network: Type[nn,Module],
+        network: Type[nn.Module],
         optimizer: Union[OmegaConf, Dict],
         lr_scheduler: Union[OmegaConf, Dict],
         criterion: Union[OmegaConf, Dict],
@@ -40,14 +40,14 @@ class LitBaseModel(pl.LightningModule):
         args = {} or criterion.args
         return getattr(nn, criterion.type)(**args)
 
-    def _configure_optimizer(self) -> type:
+    def _configure_optimizer(self) -> torch.optim.Optimizer:
         """Configures the optimizer."""
         args = {} or self._optimizer.args
         if self._optimizer.type == "MADGRAD":
             optimizer_class = madgrad.MADGRAD
         else:
             optimizer_class = getattr(torch.optim, self._optimizer.type)
-        return optimizer_class(parameters=self.parameters(), **args)
+        return optimizer_class(params=self.parameters(), **args)
 
     def _configure_lr_scheduler(self) -> Dict[str, Any]:
         """Configures the lr scheduler."""
diff --git a/text_recognizer/models/transformer.py b/text_recognizer/models/transformer.py
index 3625ab2..983e274 100644
--- a/text_recognizer/models/transformer.py
+++ b/text_recognizer/models/transformer.py
@@ -19,16 +19,14 @@ class LitTransformerModel(LitBaseModel):
 
     def __init__(
         self,
-        network: Type[nn,Module],
+        network: Type[nn, Module],
         optimizer: Union[OmegaConf, Dict],
         lr_scheduler: Union[OmegaConf, Dict],
         criterion: Union[OmegaConf, Dict],
         monitor: str = "val_loss",
         mapping: Optional[List[str]] = None,
     ) -> None:
-        super().__init__(
-            network, optimizer, lr_scheduler, criterion, monitor
-        )
+        super().__init__(network, optimizer, lr_scheduler, criterion, monitor)
 
         self.mapping, ignore_tokens = self.configure_mapping(mapping)
         self.val_cer = CharacterErrorRate(ignore_tokens)