def test_migrate_model_checkpoint_save_on_train_epoch_end_default(save_on_train_epoch_end):
    """Test that the 'save_on_train_epoch_end' part of the ModelCheckpoint state key gets removed."""
    legacy_state_key = (
        f"ModelCheckpoint{{'monitor': None, 'mode': 'min', 'every_n_train_steps': 0, 'every_n_epochs': 1,"
        f" 'train_time_interval': None, 'save_on_train_epoch_end': {save_on_train_epoch_end}}}"
    )
    new_state_key = (
        "ModelCheckpoint{'monitor': None, 'mode': 'min', 'every_n_train_steps': 0, 'every_n_epochs': 1,"
        " 'train_time_interval': None}"
    )
    old_checkpoint = {"callbacks": {legacy_state_key: {"dummy": 0}}, "global_step": 0, "epoch": 1}
    _set_version(old_checkpoint, "1.8.9")  # pretend a checkpoint prior to 1.9.0
    updated_checkpoint, _ = migrate_checkpoint(old_checkpoint, target_version="1.9.0")
    assert updated_checkpoint["callbacks"] == {new_state_key: {"dummy": 0}}  # None -> None
Test that the 'save_on_train_epoch_end' part of the ModelCheckpoint state key gets removed.
tests/tests_pytorch/utilities/migration/test_migration.py