1 files changed, 3 insertions, 3 deletions

diff --git a/training/lr.py b/training/lr.py
index 37588b6..a3144ba 100644
--- a/training/lr.py
+++ b/training/lr.py
@@ -1,6 +1,6 @@
 import math
 import copy
-from typing import Callable
+from typing import Callable, Any, Tuple, Union
 from functools import partial
 
 import matplotlib.pyplot as plt
@@ -24,7 +24,7 @@ class LRFinder():
         optimizer,
         train_dataloader,
         val_dataloader,
-        loss_fn,
+        loss_fn: Union[Callable[[Any], Tuple[Any, Any, int]], Callable[[Any, bool], Tuple[Any, Any, int]]],
         on_train: Callable[[], None] = noop,
         on_eval: Callable[[], None] = noop
     ):
@@ -108,7 +108,7 @@ class LRFinder():
                     if step >= num_val_batches:
                         break
 
-                    loss, acc, bsz = self.loss_fn(batch)
+                    loss, acc, bsz = self.loss_fn(batch, True)
                     avg_loss.update(loss.detach_(), bsz)
                     avg_acc.update(acc.detach_(), bsz)