Fix resetting epoch loop restarting flag in LearningRateFinder (#19819)

Lightning-AI · Jun 7, 2024 · 06ea3a0 · 06ea3a0
1 parent 5fa32d9
commit 06ea3a0
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 0 deletions.
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
@@ -63,6 +63,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed an issue causing ValueError for certain object such as TorchMetrics when dumping hyperparameters to YAML ([#19804](https://github.com/Lightning-AI/pytorch-lightning/pull/19804))
 
 
+- Fixed resetting `epoch_loop.restarting` to avoid full validation run after `LearningRateFinder` ([#19818](https://github.com/Lightning-AI/pytorch-lightning/issues/19818))
+
+
 ## [2.2.2] - 2024-04-11
 
 ### Fixed

diff --git a/src/lightning/pytorch/tuner/lr_finder.py b/src/lightning/pytorch/tuner/lr_finder.py
@@ -301,6 +301,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: