Repository: Lightning-AI/pytorch-lightning
File: src/lightning/pytorch/tuner/lr_finder.py (Python)

Code before:
    lr_finder_finished = False
    try:
        # Configure optimizer and scheduler
        lr_finder._exchange_scheduler(trainer)

        # Fit, lr & loss logged in callback
        _try_loop_run(trainer, params)

        # Prompt if we stopped early
        if trainer.global_step != num_training + start_steps:
            log.info(f"LR finder stopped early after {trainer.global_step} steps due to diverging loss.")

        # Transfer results from callback to lr finder object
        lr_finder.results.update({"lr": trainer.callbacks[0].lrs, "loss": trainer.callbacks[0].losses})
        lr_finder._total_batch_idx = trainer.fit_loop.total_batch_idx  # for debug purpose

        __lr_finder_restore_params(trainer, params)

        if trainer.progress_bar_callback:
            trainer.progress_bar_callback.enable()

        # Update results across ranks
        lr_finder.results = trainer.strategy.broadcast(lr_finder.results)
        lr_finder_finished = True
    except Exception as e:
        raise e
    finally:
        # Restore initial state of model (this will also restore the original optimizer state)
        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
        trainer.fit_loop._combined_loader = None  # reset data fetcher to avoid issues with the next fit
        trainer.fit_loop.setup_data()

    # Apply LR suggestion after restoring so it persists for the real training run
    # When used as a callback, the suggestion would otherwise be lost due to checkpoint restore
    if update_attr and lr_finder_finished:
        lr = lr_finder.suggestion()
        if lr is not None:
            # update the attribute on the LightningModule (e.g., lr or learning_rate)
            lightning_setattr(model, attr_name, lr)
            # also update the currently active optimizer(s) so training continues with the suggested LR
            for opt in trainer.optimizers or []:
                for pg in opt.param_groups:
                    pg["lr"] = lr
            log.info(f"Learning rate set to {lr}")
    return lr_finder



Review comment (suggestion):
```suggestion
    except Exception as ex:
        raise ex
```

Code after:
    lr_finder_finished = False
    try:
        # Configure optimizer and scheduler
        lr_finder._exchange_scheduler(trainer)

        # Fit, lr & loss logged in callback
        _try_loop_run(trainer, params)

        # Prompt if we stopped early
        if trainer.global_step != num_training + start_steps:
            log.info(f"LR finder stopped early after {trainer.global_step} steps due to diverging loss.")

        # Transfer results from callback to lr finder object
        lr_finder.results.update({"lr": trainer.callbacks[0].lrs, "loss": trainer.callbacks[0].losses})
        lr_finder._total_batch_idx = trainer.fit_loop.total_batch_idx  # for debug purpose

        __lr_finder_restore_params(trainer, params)

        if trainer.progress_bar_callback:
            trainer.progress_bar_callback.enable()

        # Update results across ranks
        lr_finder.results = trainer.strategy.broadcast(lr_finder.results)
        lr_finder_finished = True
    except Exception as ex:
        raise ex
    finally:
        # Restore initial state of model (this will also restore the original optimizer state)
        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
        trainer.fit_loop._combined_loader = None  # reset data fetcher to avoid issues with the next fit
        trainer.fit_loop.setup_data()

    # Apply LR suggestion after restoring so it persists for the real training run
    # When used as a callback, the suggestion would otherwise be lost due to checkpoint restore
    if update_attr and lr_finder_finished:
        lr = lr_finder.suggestion()
        if lr is not None:
            # update the attribute on the LightningModule (e.g., lr or learning_rate)
            lightning_setattr(model, attr_name, lr)
            # also update the currently active optimizer(s) so training continues with the suggested LR
            for opt in trainer.optimizers or []:
                for pg in opt.param_groups:
                    pg["lr"] = lr
            log.info(f"Learning rate set to {lr}")
    return lr_finder