def test_load_checkpoint_in_place(tmp_path):
    """Test that the object's state gets reloaded in-place."""
    strategy = SingleDeviceStrategy()  # surrogate class to test implementation in base class

    # objects with initial state
    saved_model = torch.nn.Linear(2, 2)
    saved_optimizer = torch.optim.Adam(saved_model.parameters(), lr=0.1)
    saved_state = {"model": saved_model, "optimizer": saved_optimizer, "int": 1, "dict": {"cocofruit": 2}}
    strategy.save_checkpoint(tmp_path / "checkpoint", state=saved_state)

    # same objects with different state
    model = torch.nn.Linear(2, 2)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.3)
    state = {"model": model, "optimizer": optimizer, "int": 10, "dict": {"cocofruit": 20}}
    assert not torch.equal(model.weight, saved_model.weight)
    assert optimizer.state_dict() != saved_optimizer.state_dict()

    remainder = strategy.load_checkpoint(tmp_path / "checkpoint", state)
    assert torch.equal(model.weight, saved_model.weight)
    assert optimizer.state_dict() == saved_optimizer.state_dict()
    assert state["int"] == saved_state["int"]
    assert state["dict"] == saved_state["dict"]
    assert not remainder

    # partial load - only model, no optimizer
    model = torch.nn.Linear(2, 2)
    state = {"model": model}
    remainder = strategy.load_checkpoint(tmp_path / "checkpoint", state)
    assert torch.equal(model.weight, saved_model.weight)
    assert list(remainder.keys()) == ["optimizer", "int", "dict"]
Test that the object's state gets reloaded in-place.
tests/tests_fabric/strategies/test_strategy.py