code: def test_elastic_checkpoint_fixed_dp(
      self, tmpdir, elastic_save, elastic_load, load_optim
  ):
      ds_config = {
          "train_batch_size": 2,
          "optimizer": {"type": "Adam"},
          "fp16": {"enabled": True, "initial_scale_power": 8},
          "zero_optimization": {
              "stage": 2,
              "elastic_checkpoint": elastic_save,
          },
      }
      hidden_dim = 10
  
      # torch 1.2.* stores raw tensor id numbers in checkpoint state which leads to
      # false positive mismatches in checkpoint state comparisons.
      # Newer torch versions store tensor ids as 0, 1, 2, ...
      expected_mismatch_keys = (
          [] if required_minimum_torch_version(1, 4) else ["params"]
      )
      models = [SimpleModel(hidden_dim) for _ in range(2)]
      model, _, _, _ = deepspeed.initialize(
          config=ds_config,
          model=models[0],
          model_parameters=models[0].parameters(),
      )
      data_loader = random_dataloader(
          model=model,
          total_samples=8,
          hidden_dim=hidden_dim,
          device=model.device,
      )
      for n, batch in enumerate(data_loader):
          loss = model(batch[0], batch[1])
          model.backward(loss)
          model.step()
      if load_optim:
          torch.save(
              model.optimizer.optimizer.state_dict(),
              os.path.join(tmpdir, "opt-state-dict"),
          )
      model.save_checkpoint(tmpdir)
  
      ds_config["zero_optimization"][
          "elastic_checkpoint"
      ] = elastic_load
      model, _, _, _ = deepspeed.initialize(
          config=ds_config,
          model=models[1],
          model_parameters=models[1].parameters(),
      )
      model.load_checkpoint(tmpdir, load_optimizer_states=load_optim)
  
      if load_optim:
          saved_sd = torch.load(os.path.join(tmpdir, "opt-state-dict"))
          curr_sd = model.optimizer.optimizer.state_dict()
          for curr_param_group, saved_param_group in zip(
              curr_sd["param_groups"], saved_sd["param_groups"]
          ):
              compare_state_dicts(
                  curr_param_group,
                  saved_param_group,
                  expected_mismatch_keys,
              )
  
      data_loader = random_dataloader(
          model=model,
          total_samples=8,
          hidden_dim=hidden_dim,
          device=model.device,
      )
      for n, batch in enumerate(data_loader):
          loss = model(batch[0], batch[1])
          model.backward(loss)
          model.step()
  
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 TestZeROElasticCheckpoint obj's underlying member method code with example usage: obj.test_elastic_checkpoint_fixed_dp(tmpdir, elastic_save, elastic_load, load_optim) without return types