From 2ce975882d9d60708725b2355b1524a5718642ce Mon Sep 17 00:00:00 2001
From: Bas Krahmer <baskrahmer@gmail.com>
Date: Thu, 18 May 2023 14:46:44 +0200
Subject: [PATCH] Bugfix: LR finder max val batches (#17636)

---
 src/lightning/pytorch/tuner/lr_finder.py    |  1 +
 tests/tests_pytorch/tuner/test_lr_finder.py | 26 +++++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/lightning/pytorch/tuner/lr_finder.py b/src/lightning/pytorch/tuner/lr_finder.py
index f776ac091d8e5..745f9e51bfaff 100644
--- a/src/lightning/pytorch/tuner/lr_finder.py
+++ b/src/lightning/pytorch/tuner/lr_finder.py
@@ -297,6 +297,7 @@ def _lr_find(
     trainer._checkpoint_connector.restore(ckpt_path)
     trainer.strategy.remove_checkpoint(ckpt_path)
     trainer.fit_loop.restarting = False  # reset restarting flag as checkpoint restoring sets it to True
+    trainer.fit_loop.epoch_loop.val_loop._combined_loader = None
 
     return lr_finder
 
diff --git a/tests/tests_pytorch/tuner/test_lr_finder.py b/tests/tests_pytorch/tuner/test_lr_finder.py
index 329334f02bcca..bc8e529def2cb 100644
--- a/tests/tests_pytorch/tuner/test_lr_finder.py
+++ b/tests/tests_pytorch/tuner/test_lr_finder.py
@@ -477,3 +477,29 @@ def test_lr_finder_with_ddp(tmpdir):
     lr = trainer.strategy.broadcast(lr)
     assert trainer.lightning_module.lr == lr
     assert lr != init_lr
+
+
+def test_lr_finder_callback_val_batches(tmpdir):
+    """Test that `LearningRateFinder` does not limit the number of val batches during training."""
+
+    class CustomBoringModel(BoringModel):
+        def __init__(self, lr):
+            super().__init__()
+            self.lr = lr
+
+        def configure_optimizers(self):
+            return torch.optim.SGD(self.parameters(), lr=self.lr)
+
+    num_lr_tuner_training_steps = 5
+    model = CustomBoringModel(0.1)
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        num_sanity_val_steps=0,
+        max_epochs=1,
+        enable_model_summary=False,
+        callbacks=[LearningRateFinder(num_training_steps=num_lr_tuner_training_steps)],
+    )
+    trainer.fit(model)
+
+    assert trainer.num_val_batches[0] == len(trainer.val_dataloaders)
+    assert trainer.num_val_batches[0] != num_lr_tuner_training_steps