code: def test_immediate_save_load(self, tmpdir, zero_stage):
      config_dict = {
          "train_batch_size": 4,
          "optimizer": {"type": "Adam"},
          "fp16": {"enabled": True, "initial_scale_power": 8},
          "zero_optimization": {
              "stage": zero_stage,
          },
      }
      hidden_dim = 10
      model = SimpleModel(hidden_dim)
  
      ds_model = create_deepspeed_model(
          config_dict=config_dict, model=model, base_optimizer=None
      )
      ds_model.save_checkpoint(tmpdir)
      ds_model.load_checkpoint(
          tmpdir,
          load_optimizer_states=False,
          load_lr_scheduler_states=False,
          load_module_only=False,
      )
  
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 1 for calling the TestZeROSaveLoadEdgeCase obj's underlying member method code with example usage: obj.test_immediate_save_load(tmpdir, zero_stage) without return types