code: def test_load_immediate_save(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)
  
      # 1. pretrain a model and save it
      dtype = torch.half
      ds_model = create_deepspeed_model(
          config_dict=config_dict, model=model, base_optimizer=None
      )
      data_loader = random_dataloader(
          model=ds_model,
          total_samples=1,
          hidden_dim=hidden_dim,
          device=ds_model.device,
          dtype=dtype,
      )
      for _, batch in enumerate(data_loader):
          loss = ds_model(batch[0], batch[1])
          ds_model.backward(loss)
          ds_model.step()
      ds_model.save_checkpoint(tmpdir)
  
      # 2. load and immediately save a model with a fresh ds engine
      ds_model = create_deepspeed_model(
          config_dict=config_dict, model=model, base_optimizer=None
      )
      ds_model.load_checkpoint(
          tmpdir,
          load_optimizer_states=False,
          load_lr_scheduler_states=False,
          load_module_only=False,
      )
      ds_model.save_checkpoint(tmpdir)
  
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 2 for calling the TestZeROSaveLoadEdgeCase obj's underlying member method code with example usage: obj.test_load_immediate_save(tmpdir, zero_stage) without return types