def test_save_checkpoint_convert_stateful_objects(tmp_path):
    """Test that when modules and optimizers are at the top-level in the state, their `state_dict()` gets used."""
    strategy = SingleDeviceStrategy()  # surrogate class to test implementation in base class
    save_checkpoint_mock = Mock()
    strategy.checkpoint_io.save_checkpoint = save_checkpoint_mock

    model = torch.nn.Linear(3, 3)
    optimizer = torch.optim.Adam(model.parameters())
    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, 0.99)

    anything = {"cocofruit": 1}
    state = {"model": model, "optimizer": optimizer, "scheduler": scheduler, "anything": anything}
    expected = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "anything": anything,
    }
    strategy.save_checkpoint(tmp_path, state)
    assert save_checkpoint_mock.call_args[1]["checkpoint"].keys() == expected.keys()
    saved_model_state = save_checkpoint_mock.call_args[1]["checkpoint"]["model"]
    assert all(torch.equal(p0, p1) for p0, p1 in zip(saved_model_state.values(), expected["model"].values()))
    assert save_checkpoint_mock.call_args[1]["checkpoint"]["optimizer"] == expected["optimizer"]
    assert save_checkpoint_mock.call_args[1]["checkpoint"]["scheduler"] == expected["scheduler"]
    assert save_checkpoint_mock.call_args[1]["checkpoint"]["anything"] == expected["anything"]
Test that when modules and optimizers are at the top-level in the state, their `state_dict()` gets used.
tests/tests_fabric/strategies/test_strategy.py