def test_model_checkpoint_only_weights(tmp_path):
    """Tests use case where ModelCheckpoint is configured to save only model weights, and user tries to load checkpoint
    to resume training."""
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmp_path,
        max_epochs=1,
        limit_train_batches=1,
        limit_val_batches=1,
        callbacks=[ModelCheckpoint(dirpath=tmp_path, save_weights_only=True)],
    )
    # fit model
    trainer.fit(model)
    # training complete
    assert trainer.state.finished, f"Training failed with {trainer.state}"

    checkpoint_path = trainer.checkpoint_callback.best_model_path

    # assert saved checkpoint has no trainer data
    checkpoint = torch.load(checkpoint_path, weights_only=True)
    assert "optimizer_states" not in checkpoint, "checkpoint should contain only model weights"
    assert "lr_schedulers" not in checkpoint, "checkpoint should contain only model weights"

    # assert loading model works when checkpoint has only weights
    assert BoringModel.load_from_checkpoint(checkpoint_path=checkpoint_path)

    # directly save model
    new_weights_path = os.path.join(tmp_path, "save_test.ckpt")
    trainer.save_checkpoint(new_weights_path, weights_only=True)
    # assert saved checkpoint has no trainer data
    checkpoint = torch.load(new_weights_path, weights_only=True)
    assert "optimizer_states" not in checkpoint, "checkpoint should contain only model weights"
    assert "lr_schedulers" not in checkpoint, "checkpoint should contain only model weights"

    # assert restoring train state fails
    with pytest.raises(KeyError, match="checkpoint contains only the model"):
        trainer._checkpoint_connector.restore(new_weights_path)
Tests use case where ModelCheckpoint is configured to save only model weights, and user tries to load checkpoint
    to resume training.
tests/tests_pytorch/trainer/test_trainer.py