code: 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"]
      )
  
imports: from unittest import mock
  from unittest.mock import Mock, PropertyMock
  import pytest
  import torch
  from lightning.fabric.plugins import DoublePrecision, HalfPrecision, Precision
  from lightning.fabric.strategies import SingleDeviceStrategy
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
  from lightning.fabric.utilities.types import _Stateful
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 3 for calling a global function with example usage: test_save_checkpoint_convert_stateful_objects(tmp_path) without return types