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

Code before:
    # Save initial model, that is loaded after batch size is found
    ckpt_path = os.path.join(trainer.default_root_dir, f".scale_batch_size_{uuid.uuid4()}.ckpt")
    trainer.save_checkpoint(ckpt_path)

    # Arguments we adjust during the batch size finder, save for restoring
    params = __scale_batch_dump_params(trainer)

    # Set to values that are required by the algorithm
    __scale_batch_reset_params(trainer, steps_per_trial)

    if trainer.progress_bar_callback:
        trainer.progress_bar_callback.disable()

    try:
        new_size, _ = _adjust_batch_size(trainer, batch_arg_name, value=init_val)

        if mode == "power":
            new_size = _run_power_scaling(trainer, new_size, batch_arg_name, max_trials, params)
        elif mode == "binsearch":
            new_size = _run_binary_scaling(trainer, new_size, batch_arg_name, max_trials, params)

        garbage_collection_cuda()

        log.info(f"Finished batch size finder, will continue with full run using batch size {new_size}")
    except Exception as e:
        raise e
    finally:
        __scale_batch_restore_params(trainer, params)

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

        trainer._checkpoint_connector.restore(ckpt_path)
        trainer.strategy.remove_checkpoint(ckpt_path)

    return new_size


def __scale_batch_dump_params(trainer: "pl.Trainer") -> dict[str, Any]:
    dumped_params = {
        "loggers": trainer.loggers,
        "callbacks": trainer.callbacks,
    }
    loop = trainer._active_loop
    assert loop is not None
    if isinstance(loop, pl.loops._FitLoop):
        dumped_params["max_steps"] = trainer.max_steps
        dumped_params["limit_train_batches"] = trainer.limit_train_batches
        dumped_params["limit_val_batches"] = trainer.limit_val_batches
    elif isinstance(loop, pl.loops._EvaluationLoop):
        stage = trainer.state.stage

Review comment (suggestion):
```suggestion
    except Exception as ex:
        raise ex
```
pls not a single letter :)

Code after:
    # Save initial model, that is loaded after batch size is found
    ckpt_path = os.path.join(trainer.default_root_dir, f".scale_batch_size_{uuid.uuid4()}.ckpt")
    trainer.save_checkpoint(ckpt_path)

    # Arguments we adjust during the batch size finder, save for restoring
    params = __scale_batch_dump_params(trainer)

    # Set to values that are required by the algorithm
    __scale_batch_reset_params(trainer, steps_per_trial)

    if trainer.progress_bar_callback:
        trainer.progress_bar_callback.disable()

    try:
        new_size, _ = _adjust_batch_size(trainer, batch_arg_name, value=init_val)

        if mode == "power":
            new_size = _run_power_scaling(trainer, new_size, batch_arg_name, max_trials, params)
        elif mode == "binsearch":
            new_size = _run_binary_scaling(trainer, new_size, batch_arg_name, max_trials, params)

        garbage_collection_cuda()

        log.info(f"Finished batch size finder, will continue with full run using batch size {new_size}")
    except Exception as ex:
        raise ex
    finally:
        __scale_batch_restore_params(trainer, params)

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

        trainer._checkpoint_connector.restore(ckpt_path)
        trainer.strategy.remove_checkpoint(ckpt_path)

    return new_size


def __scale_batch_dump_params(trainer: "pl.Trainer") -> dict[str, Any]:
    dumped_params = {
        "loggers": trainer.loggers,
        "callbacks": trainer.callbacks,
    }
    loop = trainer._active_loop
    assert loop is not None
    if isinstance(loop, pl.loops._FitLoop):
        dumped_params["max_steps"] = trainer.max_steps
        dumped_params["limit_train_batches"] = trainer.limit_train_batches
        dumped_params["limit_val_batches"] = trainer.limit_val_batches
    elif isinstance(loop, pl.loops._EvaluationLoop):
        stage = trainer.state.stage