code: def run(self, class_tmpdir, elastic_save, 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_save,
          },
      }
      hidden_dim = 10
      model = SimpleModel(hidden_dim)
  
      model, _, _, _ = deepspeed.initialize(
          config=ds_config,
          model=model,
          model_parameters=model.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(class_tmpdir, "opt-state-dict"),
          )
      model.save_checkpoint(class_tmpdir)
  
imports: import deepspeed
  from types import SimpleNamespace
  from deepspeed.ops.op_builder import CPUAdamBuilder
  from deepspeed.checkpoint.utils import clone_tensors_for_torch_save, get_model_ckpt_name_for_rank
  from deepspeed.accelerator import get_accelerator
  from deepspeed.runtime.utils import required_torch_version
  from unit.common import DistributedTest, DistributedFixture
  from unit.simple_model import *
  from unit.checkpoint.common import *
  import pytest
label_desc: Use image node_id 1 for calling the ws4_model_checkpoint obj's underlying member method code with example usage: obj.run(class_tmpdir, elastic_save, load_optim) without return types