diff options
Diffstat (limited to 'text_recognizer/models')
| -rw-r--r-- | text_recognizer/models/base.py | 6 | ||||
| -rw-r--r-- | text_recognizer/models/transformer.py | 6 |
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) |