fix: lr_finder epoch loop restarting flag (Lightning-AI#19818)

clumsy · Apr 25, 2024 · 89cd929 · 89cd929
1 parent b9680a3
commit 89cd929
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 0 deletions.
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
@@ -16,6 +16,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `on_exception` hook to `LightningDataModule` ([#19601](https://github.com/Lightning-AI/pytorch-lightning/pull/19601))
 
+- Set `epoch_loop.restarting` to `False` to avoid full validation run after `LearningRateFinder` [#19818](https://github.com/Lightning-AI/pytorch-lightning/issues/19818))
+
 -
 
 ### Changed

diff --git a/src/lightning/pytorch/tuner/lr_finder.py b/src/lightning/pytorch/tuner/lr_finder.py
@@ -302,6 +302,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.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
@@ -434,6 +434,7 @@ def lr_find(self, trainer, pl_module) -> None:
             super().lr_find(trainer, pl_module)
             pl_module._expected_max_steps = None
             assert not trainer.fit_loop.restarting
+            assert not trainer.fit_loop.epoch_loop.restarting
 
         def on_train_epoch_start(self, trainer, pl_module):
             if trainer.current_epoch in self.milestones or trainer.current_epoch == 0: