def test_ckpt_every_n_train_steps(tmp_path):
    """Tests that the checkpoints are saved every n training steps."""
    model = LogInTwoMethods()
    every_n_train_steps = 16
    max_epochs = 2
    epoch_length = 64
    checkpoint_callback = ModelCheckpoint(
        filename="{step}",
        every_n_epochs=0,
        every_n_train_steps=every_n_train_steps,
        dirpath=tmp_path,
        save_top_k=-1,
        save_last=False,
    )
    trainer = Trainer(
        default_root_dir=tmp_path,
        max_epochs=2,
        enable_progress_bar=False,
        callbacks=[checkpoint_callback],
        logger=False,
    )

    trainer.fit(model)
    expected = [
        f"step={i}.ckpt" for i in range(every_n_train_steps, max_epochs * epoch_length + 1, every_n_train_steps)
    ]
    assert set(os.listdir(tmp_path)) == set(expected)
Tests that the checkpoints are saved every n training steps.
tests/tests_pytorch/checkpointing/test_model_checkpoint.py