2 files changed, 26 insertions, 27 deletions

diff --git a/training/lr.py b/training/lr.py
index c765150..68e0f72 100644
--- a/training/lr.py
+++ b/training/lr.py
@@ -1,5 +1,5 @@
 import math
-import copy
+from contextlib import _GeneratorContextManager, nullcontext
 from typing import Callable, Any, Tuple, Union
 from functools import partial
 
@@ -25,9 +25,9 @@ class LRFinder():
         train_dataloader,
         val_dataloader,
         loss_fn: Union[Callable[[int, Any], Tuple[Any, Any, int]], Callable[[int, Any, bool], Tuple[Any, Any, int]]],
-        on_train: Callable[[], None] = noop,
+        on_train: Callable[[], _GeneratorContextManager] = nullcontext,
         on_clip: Callable[[], None] = noop,
-        on_eval: Callable[[], None] = noop
+        on_eval: Callable[[], _GeneratorContextManager] = nullcontext
     ):
         self.accelerator = accelerator
         self.model = model
@@ -51,7 +51,6 @@ class LRFinder():
         num_train_batches: int = 1,
         num_val_batches: int = math.inf,
         smooth_f: float = 0.05,
-        diverge_th: int = 5,
     ):
         best_loss = None
         best_acc = None
@@ -84,40 +83,40 @@ class LRFinder():
             avg_acc = AverageMeter()
 
             self.model.train()
-            self.on_train()
 
-            for step, batch in enumerate(self.train_dataloader):
-                if step >= num_train_batches:
-                    break
+            with self.on_train():
+                for step, batch in enumerate(self.train_dataloader):
+                    if step >= num_train_batches:
+                        break
 
-                with self.accelerator.accumulate(self.model):
-                    loss, acc, bsz = self.loss_fn(step, batch)
+                    with self.accelerator.accumulate(self.model):
+                        loss, acc, bsz = self.loss_fn(step, batch)
 
-                    self.accelerator.backward(loss)
+                        self.accelerator.backward(loss)
 
-                    if self.accelerator.sync_gradients:
-                        self.on_clip()
+                        if self.accelerator.sync_gradients:
+                            self.on_clip()
 
-                    self.optimizer.step()
-                    lr_scheduler.step()
-                    self.optimizer.zero_grad(set_to_none=True)
+                        self.optimizer.step()
+                        lr_scheduler.step()
+                        self.optimizer.zero_grad(set_to_none=True)
 
-                if self.accelerator.sync_gradients:
-                    progress_bar.update(1)
+                    if self.accelerator.sync_gradients:
+                        progress_bar.update(1)
 
             self.model.eval()
-            self.on_eval()
 
             with torch.inference_mode():
-                for step, batch in enumerate(self.val_dataloader):
-                    if step >= num_val_batches:
-                        break
+                with self.on_eval():
+                    for step, batch in enumerate(self.val_dataloader):
+                        if step >= num_val_batches:
+                            break
 
-                    loss, acc, bsz = self.loss_fn(step, batch, True)
-                    avg_loss.update(loss.detach_(), bsz)
-                    avg_acc.update(acc.detach_(), bsz)
+                        loss, acc, bsz = self.loss_fn(step, batch, True)
+                        avg_loss.update(loss.detach_(), bsz)
+                        avg_acc.update(acc.detach_(), bsz)
 
-                    progress_bar.update(1)
+                        progress_bar.update(1)
 
             loss = avg_loss.avg.item()
             acc = avg_acc.avg.item()
diff --git a/training/util.py b/training/util.py
index 6f1e85a..bed7111 100644
--- a/training/util.py
+++ b/training/util.py
@@ -262,7 +262,7 @@ class EMAModel:
                 raise ValueError("collected_params and shadow_params must have the same length")
 
     @contextmanager
-    def apply_temporary(self, parameters):
+    def apply_temporary(self, parameters: Iterable[torch.nn.Parameter]):
         try:
             parameters = list(parameters)
             original_params = [p.clone() for p in parameters]