def test_fit_loop_reset(tmp_path):
    """Test that the reset logic in fit- and epoch loop is aware of whether the loop is restarting from a completed
    loop or from a mid-epoch checkpoint."""

    # generate checkpoints at end of epoch and mid-epoch
    model = BoringModel()
    checkpoint_callback = ModelCheckpoint(
        dirpath=tmp_path,
        every_n_train_steps=2,
        save_top_k=-1,
    )
    trainer = Trainer(
        default_root_dir=tmp_path,
        limit_train_batches=4,
        max_epochs=2,
        callbacks=[checkpoint_callback],
        logger=False,
        enable_model_summary=False,
    )
    trainer.fit(model)

    # reset state loaded from a checkpoint from mid-epoch
    mid_epoch_ckpt = torch.load(str(tmp_path / "epoch=0-step=2.ckpt"), weights_only=True)
    fit_loop = trainer.fit_loop
    epoch_loop = fit_loop.epoch_loop
    optimizer_loop = epoch_loop.automatic_optimization
    assert not fit_loop.restarting
    assert not epoch_loop.restarting
    assert not optimizer_loop.restarting

    # we load exactly what was saved - no reset yet
    fit_loop.load_state_dict(mid_epoch_ckpt["loops"]["fit_loop"])

    assert fit_loop.restarting
    assert fit_loop.epoch_progress.total.ready == 1
    assert fit_loop.epoch_progress.total.completed == 0  # the checkpoint was saved mid epoch
    assert fit_loop.epoch_progress.current.ready == 1
    assert fit_loop.epoch_progress.current.completed == 0

    assert epoch_loop.batch_progress.total.ready == 2
    assert epoch_loop.batch_progress.total.processed == 2
    assert epoch_loop.batch_progress.total.completed == 1  # the checkpoint was saved on train_batch_end
    assert epoch_loop.batch_progress.current.ready == 2  # currents get set to the completed value
    assert epoch_loop.batch_progress.current.processed == 2
    assert epoch_loop.batch_progress.current.completed == 1

    fit_loop.reset()
    epoch_loop.reset()

    # resetting from a mid-of-epoch checkpoint SHOULD NOT reset the current counters to 0
    assert fit_loop.restarting
    assert fit_loop.epoch_progress.total.ready == 1
    assert fit_loop.epoch_progress.total.completed == 0  # the checkpoint was saved mid epoch
    assert fit_loop.epoch_progress.current.ready == 1
    assert fit_loop.epoch_progress.current.completed == 0

    # however it should increment completed batch progress, since it was saved immediately prior
    assert epoch_loop.restarting
    assert epoch_loop.batch_progress.total.ready == 2
    assert epoch_loop.batch_progress.total.processed == 2
    assert epoch_loop.batch_progress.total.completed == 2
    assert epoch_loop.batch_progress.current.ready == 2
    assert epoch_loop.batch_progress.current.processed == 2
    assert epoch_loop.batch_progress.current.completed == 2

    assert optimizer_loop.restarting

    # reset state loaded from a checkpoint from the end of an epoch
    end_of_epoch_ckpt = torch.load(str(tmp_path / "epoch=0-step=4.ckpt"), weights_only=True)
    fit_loop = trainer.fit_loop
    epoch_loop = fit_loop.epoch_loop
    fit_loop.restarting = False
    epoch_loop.restarting = False
    optimizer_loop.restarting = False

    # we load exactly what was saved - no reset yet
    fit_loop.load_state_dict(end_of_epoch_ckpt["loops"]["fit_loop"])

    assert fit_loop.restarting
    assert fit_loop.epoch_progress.total.ready == 1
    assert fit_loop.epoch_progress.total.completed == 0
    assert fit_loop.epoch_progress.current.ready == 1
    assert fit_loop.epoch_progress.current.completed == 0

    # resetting from a end-of-epoch checkpoint SHOULD reset the current counters to 0
    fit_loop.reset()
    epoch_loop.reset()

    # resetting from a mid-of-epoch checkpoint SHOULD NOT reset the current counters to 0
    # since we are restarting at the end of epoch, we need to see `completed` being updated after reset
    assert fit_loop.restarting
    assert fit_loop.epoch_progress.total.ready == 1
    assert fit_loop.epoch_progress.total.completed == 1
    assert fit_loop.epoch_progress.current.ready == 1
    assert fit_loop.epoch_progress.current.completed == 1

    # however it should increment completed batch progress, since it was saved immediately prior
    assert epoch_loop.restarting
    assert epoch_loop.batch_progress.total.ready == 4
    assert epoch_loop.batch_progress.total.processed == 4
    assert epoch_loop.batch_progress.total.completed == 4
    assert epoch_loop.batch_progress.current.ready == 0
    assert epoch_loop.batch_progress.current.processed == 0
    assert epoch_loop.batch_progress.current.completed == 0
Test that the reset logic in fit- and epoch loop is aware of whether the loop is restarting from a completed
    loop or from a mid-epoch checkpoint.
tests/tests_pytorch/loops/test_loops.py