From b89953bea7dfe6c92164888a66d05bc7d987ef71 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Wed, 11 Jan 2023 22:29:28 +0100 Subject: Fix --- training/lr.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'training') 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(): val_dataloader, loss_fn: Union[Callable[[int, Any], Tuple[Any, Any, int]], Callable[[int, Any, bool], Tuple[Any, Any, int]]], on_train: Callable[[], _GeneratorContextManager] = nullcontext, - on_clip: Callable[[], None] = noop, + on_clip: Callable[[float], None] = noop, on_eval: Callable[[], _GeneratorContextManager] = nullcontext ): self.accelerator = accelerator @@ -95,7 +95,7 @@ class LRFinder(): self.accelerator.backward(loss) if self.accelerator.sync_gradients: - self.on_clip() + self.on_clip(lr_scheduler.get_last_lr()[0]) self.optimizer.step() lr_scheduler.step() -- cgit v1.2.3-54-g00ecf