diff --git a/src/mednet/libs/common/engine/callbacks.py b/src/mednet/libs/common/engine/callbacks.py
index 4ef56073a3731dc6b81aa1d4ef095f18e48618d1..0f17f098d69d723fb89ccb12e152d0ea2e58ca88 100644
--- a/src/mednet/libs/common/engine/callbacks.py
+++ b/src/mednet/libs/common/engine/callbacks.py
@@ -143,12 +143,11 @@ class LoggingCallback(lightning.pytorch.Callback):
         self._to_log["epoch-duration-seconds/train"] = epoch_time
 
         if len(pl_module.trainer.lr_scheduler_configs) == 0:
-            self._to_log["learning-rate"] = pl_module.optimizers().defaults["lr"]
+            self._to_log["learning-rate"] = pl_module.optimizers().defaults["lr"]  # type: ignore
         else:
-            __import__("pdb").set_trace()
             self._to_log["learning-rate"] = pl_module.trainer.lr_scheduler_configs[
                 0
-            ].scheduler.optimizer.param_groups[0]["lr"]  # type: ignore
+            ].scheduler.get_last_lr()  # type: ignore
 
         overall_cycle_time = time.time() - self._start_training_epoch_time
         self._to_log["cycle-time-seconds/train"] = overall_cycle_time