Fix restarting attribute for lr finder

Author: justusschock

src/pytorch_lightning/tuner/lr_finder.py

@@ -270,6 +270,7 @@ def lr_find(
     # Restore initial state of model
     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
 
     return lr_finder
 

tests/tests_pytorch/tuner/test_lr_finder.py

@@ -441,6 +441,43 @@ def test_if_lr_finder_callback_already_configured():
         trainer.tune(model)
 
 
+def test_lr_finder_callback_restarting(tmpdir):
+    """Test that `LearningRateFinder` does not set restarting=True when loading checkpoint."""
+
+    class MyBoringModel(BoringModel):
+        def __init__(self):
+            super().__init__()
+            self.learning_rate = 0.123
+
+        def configure_optimizers(self):
+            return torch.optim.SGD(self.parameters(), lr=self.learning_rate)
+
+    class CustomLearningRateFinder(LearningRateFinder):
+        milestones = (1,)
+
+        def lr_find(self, trainer, pl_module) -> None:
+            super().lr_find(trainer, pl_module)
+            assert not trainer.fit_loop.restarting
+
+        def on_train_epoch_start(self, trainer, pl_module):
+            if trainer.current_epoch in self.milestones or trainer.current_epoch == 0:
+                self.lr_find(trainer, pl_module)
+
+    model = MyBoringModel()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=3,
+        callbacks=[CustomLearningRateFinder(early_stop_threshold=None, update_attr=True)],
+        limit_train_batches=10,
+        limit_val_batches=0,
+        limit_test_batches=00,
+        num_sanity_val_steps=0,
+        enable_model_summary=False,
+    )
+
+    trainer.fit(model)
+
+
 @mock.patch.dict(os.environ, os.environ.copy(), clear=True)
 @RunIf(standalone=True)
 def test_lr_finder_with_ddp(tmpdir):