code: def test_hybrid_optimizer_state(self, tmpdir, zero_stage):
      config_dict = {
          "train_micro_batch_size_per_gpu": 2,
          "gradient_accumulation_steps": 2,
          "steps_per_print": 1,
          "zero_optimization": {"stage": zero_stage},
          "zero_allow_untested_optimizer": True,
          "fp16": {"enabled": True, "initial_scale_power": 8},
      }
      hidden_dim = 10
      models = [SimpleModel(hidden_dim=hidden_dim) for _ in range(2)]
      optimizers = [
          HybridStateOptimizer(model.parameters()) for model in models
      ]
  
      checkpoint_correctness_verification(
          config_dict,
          models=models,
          base_optimizers=optimizers,
          hidden_dim=hidden_dim,
          tmpdir=tmpdir,
          load_optimizer_states=True,
      )
  
imports: import deepspeed
  from deepspeed.ops.op_builder import CPUAdamBuilder
  from unit.common import DistributedTest, DistributedFixture
  from unit.simple_model import *
  from unit.util import required_minimum_torch_version
  from unit.checkpoint.common import *
  import pytest
label_desc: Use image node_id 3 for calling the TestZeROCheckpoint obj's underlying member method code with example usage: obj.test_hybrid_optimizer_state(tmpdir, zero_stage) without return types