Training improvements

1 files changed, 4 insertions, 3 deletions

diff --git a/training/lr.py b/training/lr.py
index c0e9b3f..0c5ce9e 100644
--- a/training/lr.py
+++ b/training/lr.py
@@ -90,6 +90,7 @@ class LRFinder():
             else:
                 if smooth_f > 0:
                     loss = smooth_f * loss + (1 - smooth_f) * losses[-1]
+                    acc = smooth_f * acc + (1 - smooth_f) * accs[-1]
                 if loss < best_loss:
                     best_loss = loss
                 if acc > best_acc:
@@ -132,9 +133,9 @@ class LRFinder():
         ax_loss.set_xlabel("Learning rate")
         ax_loss.set_ylabel("Loss")
 
-        # ax_acc = ax_loss.twinx()
-        # ax_acc.plot(lrs, accs, color='blue')
-        # ax_acc.set_ylabel("Accuracy")
+        ax_acc = ax_loss.twinx()
+        ax_acc.plot(lrs, accs, color='blue')
+        ax_acc.set_ylabel("Accuracy")
 
         print("LR suggestion: steepest gradient")
         min_grad_idx = None