1 files changed, 3 insertions, 3 deletions

diff --git a/training/functional.py b/training/functional.py
index bd8cbad..b9fb546 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -39,7 +39,7 @@ class TrainingCallbacks():
     on_train: Callable[[int], _GeneratorContextManager] = const(nullcontext())
     on_before_optimize: Callable[[float, int], Any] = const()
     on_after_optimize: Callable[[Any, float], None] = const()
-    on_after_epoch: Callable[[float], None] = const()
+    on_after_epoch: Callable[[], None] = const()
     on_eval: Callable[[], _GeneratorContextManager] = const(nullcontext())
     on_sample: Callable[[int], None] = const()
     on_checkpoint: Callable[[int, str], None] = const()
@@ -496,7 +496,7 @@ def train_loop(
                         "lr": lr,
                     }
                     if isDadaptation:
-                        logs["lr/d*lr"] = optimizer.param_groups[0]["d"] * optimizer.param_groups[0]["lr"]
+                        logs["lr/d*lr"] = lr = optimizer.param_groups[0]["d"] * optimizer.param_groups[0]["lr"]
                     logs.update(on_log())
 
                     local_progress_bar.set_postfix(**logs)
@@ -528,7 +528,7 @@ def train_loop(
 
             lrs.append(lr)
 
-            on_after_epoch(lr)
+            on_after_epoch()
 
             if val_dataloader is not None:
                 model.eval()