1 files changed, 26 insertions, 44 deletions
diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py
index 26cc18c..1ebb256 100644
--- a/text_recognizer/models/base.py
+++ b/text_recognizer/models/base.py
@@ -1,5 +1,5 @@
 """Base PyTorch Lightning model."""
-from typing import Any, Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, Optional, Tuple, Type
 
 import hydra
 from loguru import logger as log
@@ -47,57 +47,39 @@ class LitBase(LightningModule):
         """Optimal way to set grads to zero."""
         optimizer.zero_grad(set_to_none=True)
 
-    def _configure_optimizer(self) -> List[Type[torch.optim.Optimizer]]:
+    def _configure_optimizer(self) -> Type[torch.optim.Optimizer]:
         """Configures the optimizer."""
-        optimizers = []
-        for optimizer_config in self.optimizer_configs.values():
-            module = self
-            for m in str(optimizer_config.parameters).split("."):
-                module = getattr(module, m)
-            del optimizer_config.parameters
-            log.info(f"Instantiating optimizer <{optimizer_config._target_}>")
-            optimizers.append(
-                hydra.utils.instantiate(optimizer_config, params=module.parameters())
-            )
-        return optimizers
+        return hydra.utils.instantiate(
+            self.optimizer_config, params=self.network.parameters()
+        )
 
     def _configure_lr_schedulers(
-        self, optimizers: List[Type[torch.optim.Optimizer]]
-    ) -> List[Dict[str, Any]]:
+        self, optimizer: Type[torch.optim.Optimizer]
+    ) -> Dict[str, Any]:
         """Configures the lr scheduler."""
-        if self.lr_scheduler_configs is None:
-            return []
-        schedulers = []
-        for optimizer, lr_scheduler_config in zip(
-            optimizers, self.lr_scheduler_configs.values()
-        ):
-            # Extract non-class arguments.
-            monitor = lr_scheduler_config.monitor
-            interval = lr_scheduler_config.interval
-            del lr_scheduler_config.monitor
-            del lr_scheduler_config.interval
-
-            log.info(
-                f"Instantiating learning rate scheduler <{lr_scheduler_config._target_}>"
-            )
-            scheduler = {
-                "monitor": monitor,
-                "interval": interval,
-                "scheduler": hydra.utils.instantiate(
-                    lr_scheduler_config, optimizer=optimizer
-                ),
-            }
-            schedulers.append(scheduler)
-
-        return schedulers
+        log.info(
+            f"Instantiating learning rate scheduler <{self.lr_scheduler_config._target_}>"
+        )
+        monitor = self.lr_scheduler_config.monitor
+        interval = self.lr_scheduler_config.interval
+        del self.lr_scheduler_config.monitor
+        del self.lr_scheduler_config.interval
+
+        return {
+            "monitor": monitor,
+            "interval": interval,
+            "scheduler": hydra.utils.instantiate(
+                self.lr_scheduler_config, optimizer=optimizer
+            ),
+        }
 
     def configure_optimizers(
         self,
-    ) -> Tuple[List[Type[torch.optim.Optimizer]], List[Dict[str, Any]]]:
+    ) -> Dict[str, Any]:
         """Configures optimizer and lr scheduler."""
-        optimizers = self._configure_optimizer()
-        schedulers = self._configure_lr_schedulers(optimizers)
-        return optimizers, schedulers
+        optimizer = self._configure_optimizer()
+        scheduler = self._configure_lr_schedulers(optimizer)
+        return {"optimizer": optimizer, "scheduler": scheduler}
 
     def forward(self, data: Tensor) -> Tensor:
         """Feedforward pass."""