summaryrefslogtreecommitdiffstats
path: root/training
diff options
context:
space:
mode:
Diffstat (limited to 'training')
-rw-r--r--training/lr.py4
1 files changed, 2 insertions, 2 deletions
diff --git a/training/lr.py b/training/lr.py
index dfb1743..01f7f5e 100644
--- a/training/lr.py
+++ b/training/lr.py
@@ -26,7 +26,7 @@ class LRFinder():
26 val_dataloader, 26 val_dataloader,
27 loss_fn: Union[Callable[[int, Any], Tuple[Any, Any, int]], Callable[[int, Any, bool], Tuple[Any, Any, int]]], 27 loss_fn: Union[Callable[[int, Any], Tuple[Any, Any, int]], Callable[[int, Any, bool], Tuple[Any, Any, int]]],
28 on_train: Callable[[], _GeneratorContextManager] = nullcontext, 28 on_train: Callable[[], _GeneratorContextManager] = nullcontext,
29 on_clip: Callable[[], None] = noop, 29 on_clip: Callable[[float], None] = noop,
30 on_eval: Callable[[], _GeneratorContextManager] = nullcontext 30 on_eval: Callable[[], _GeneratorContextManager] = nullcontext
31 ): 31 ):
32 self.accelerator = accelerator 32 self.accelerator = accelerator
@@ -95,7 +95,7 @@ class LRFinder():
95 self.accelerator.backward(loss) 95 self.accelerator.backward(loss)
96 96
97 if self.accelerator.sync_gradients: 97 if self.accelerator.sync_gradients:
98 self.on_clip() 98 self.on_clip(lr_scheduler.get_last_lr()[0])
99 99
100 self.optimizer.step() 100 self.optimizer.step()
101 lr_scheduler.step() 101 lr_scheduler.step()