1 files changed, 1 insertions, 1 deletions
diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py
index 8dbdafa..70eb8dd 100644
--- a/text_recognizer/models/base.py
+++ b/text_recognizer/models/base.py
@@ -79,7 +79,7 @@ class LitBase(LightningModule):
         """Configures optimizer and lr scheduler."""
         optimizer = self._configure_optimizer()
         scheduler = self._configure_lr_schedulers(optimizer)
-        return {"optimizer": optimizer, "scheduler": scheduler}
+        return {"optimizer": optimizer, "lr_scheduler": scheduler}
 
     def forward(self, data: Tensor) -> Tensor:
         """Feedforward pass."""