code: def test_elastic_checkpoint_change_dp(
      self,
      ws4_model_checkpoint,
      class_tmpdir,
      elastic_save,
      elastic_load,
      load_optim,
  ):
      ds_config = {
          "train_batch_size": 4,
          "optimizer": {"type": "Adam"},
          "fp16": {"enabled": True, "initial_scale_power": 8},
          "zero_optimization": {
              "stage": 2,
              "elastic_checkpoint": elastic_load,
          },
      }
      hidden_dim = 10
      model = SimpleModel(hidden_dim)
  
      # Load checkpoint with dp world size = 2
      model, _, _, _ = deepspeed.initialize(
          config=ds_config,
          model=model,
          model_parameters=model.parameters(),
      )
      if load_optim:
          with pytest.raises(
              deepspeed.runtime.zero.utils.ZeRORuntimeException
          ):
              model.load_checkpoint(
                  class_tmpdir, load_optimizer_states=load_optim
              )
      else:
          model.load_checkpoint(
              class_tmpdir, load_optimizer_states=load_optim
          )
  
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 TestZeROElasticCheckpoint obj's underlying member method code with example usage: obj.test_elastic_checkpoint_change_dp(ws4_model_checkpoint, class_tmpdir, elastic_save, elastic_load, load_optim) without return types