Repository: Lightning-AI/pytorch-lightning
File: tests/tests_pytorch/models/test_hparams.py (Python)

Code before:
            callbacks=[checkpoint_callback],
            limit_train_batches=10,
            limit_val_batches=10,
            max_epochs=epochs,
            logger=False,
        )
        trainer.fit(model)
        _ = TestHydraModel.load_from_checkpoint(checkpoint_callback.best_model_path, weights_only=False)


@pytest.mark.parametrize("ignore", ["arg2", ("arg2", "arg3")])
def test_ignore_args_list_hparams(tmp_path, ignore):
    """Tests that args can be ignored in save_hyperparameters."""

    class LocalModel(BoringModel):
        def __init__(self, arg1, arg2, arg3):
            super().__init__()
            self.save_hyperparameters(ignore=ignore)

    model = LocalModel(arg1=14, arg2=90, arg3=50)

    # test proper property assignments
    assert model.hparams.arg1 == 14

    ignore_args = ignore if isinstance(ignore, (list, tuple)) else [ignore]
    for arg in ignore_args:
        assert arg not in model.hparams

    # verify we can train
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=2, overfit_batches=0.5)
    trainer.fit(model)

    # make sure the raw checkpoint saved the properties
    raw_checkpoint_path = _raw_checkpoint_path(trainer)
    raw_checkpoint = torch.load(raw_checkpoint_path, weights_only=True)
    assert LightningModule.CHECKPOINT_HYPER_PARAMS_KEY in raw_checkpoint
    assert raw_checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_KEY]["arg1"] == 14

    # verify that model loads correctly
    model = LocalModel.load_from_checkpoint(raw_checkpoint_path, arg2=123, arg3=100)
    assert model.hparams.arg1 == 14
    for arg in ignore_args:
        assert arg not in model.hparams


@pytest.mark.parametrize("ignore", ["arg2", ("arg2", "arg3")])
def test_hparams_ignore_persists_across_subclassing(tmp_path, ignore):
    """Test that hyperparameters can be ignored when `save_hyperparameters` is called in both a base class and a
    subclass, and that ignore rules defined in the subclass override hyperparameters saved by the base class."""

    class BaseBoringModel(BoringModel):

Review comment (suggestion):
The reason for this change is when ignore="arg2", the loop was iterating over the characters of the string "arg2".

Code after:
            callbacks=[checkpoint_callback],
            limit_train_batches=10,
            limit_val_batches=10,
            max_epochs=epochs,
            logger=False,
        )
        trainer.fit(model)
        _ = TestHydraModel.load_from_checkpoint(checkpoint_callback.best_model_path, weights_only=False)


@pytest.mark.parametrize("ignore", ["arg2", ("arg2", "arg3")])
def test_ignore_args_list_hparams(tmp_path, ignore):
    """Tests that args can be ignored in save_hyperparameters."""

    class LocalModel(BoringModel):
        def __init__(self, arg1, arg2, arg3):
            super().__init__()
            self.save_hyperparameters(ignore=ignore)

    model = LocalModel(arg1=14, arg2=90, arg3=50)

    # test proper property assignments
    assert model.hparams.arg1 == 14

    ignore_args = ignore if isinstance(ignore, (list, tuple)) else [ignore]
    for arg in ignore_args:
        assert arg not in model.hparams

    # verify we can train
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=2, overfit_batches=0.5)
    trainer.fit(model)

    # make sure the raw checkpoint saved the properties
    raw_checkpoint_path = _raw_checkpoint_path(trainer)
    raw_checkpoint = torch.load(raw_checkpoint_path, weights_only=True)
    assert LightningModule.CHECKPOINT_HYPER_PARAMS_KEY in raw_checkpoint
    assert raw_checkpoint[LightningModule.CHECKPOINT_HYPER_PARAMS_KEY]["arg1"] == 14

    # verify that model loads correctly
    model = LocalModel.load_from_checkpoint(raw_checkpoint_path, arg2=123, arg3=100)
    assert model.hparams.arg1 == 14
    for arg in ignore_args:
        assert arg not in model.hparams


@pytest.mark.parametrize("ignore", ["arg2", ("arg2", "arg3")])
def test_hparams_ignore_in_subclass_overrides_base(tmp_path, ignore):
    """Test that hyperparameters can be ignored when `save_hyperparameters` is called in both a base class and a
    subclass, and that ignore rules defined in the subclass override hyperparameters saved by the base class."""

    class BaseBoringModel(BoringModel):