2 files changed, 7 insertions, 1 deletions

diff --git a/training/lr.py b/training/lr.py
index 0c5ce9e..3abd2f2 100644
--- a/training/lr.py
+++ b/training/lr.py
@@ -102,6 +102,12 @@ class LRFinder():
             losses.append(loss)
             accs.append(acc)
 
+            self.accelerator.log({
+                "loss": loss,
+                "acc": acc,
+                "lr": lr,
+            }, step=epoch)
+
             progress_bar.set_postfix({
                 "loss": loss,
                 "loss/best": best_loss,
diff --git a/training/optimization.py b/training/optimization.py
index 3340544..a79944f 100644
--- a/training/optimization.py
+++ b/training/optimization.py
@@ -15,7 +15,7 @@ def get_one_cycle_schedule(
     warmup: Literal["cos", "linear"] = "cos",
     annealing: Literal["cos", "half_cos", "linear"] = "cos",
     warmup_exp: int = 1,
-    annealing_exp: int = 2,
+    annealing_exp: int = 1,
     min_lr: int = 0.04,
     mid_point: int = 0.3,
     last_epoch: int = -1