code: def test_checkpoint_shared_weights(self, tmp_path):
      config = {
          "train_micro_batch_size_per_gpu": 2,
          "zero_allow_untested_optimizer": True,
          "zero_optimization": {"stage": 2},
      }
      model = ModelWithSharedWeights()
      optimizer = torch.optim.Adam(model.parameters())
  
      deepspeed_engine, _, _, _ = deepspeed.initialize(
          config=config,
          model=model,
          optimizer=optimizer,
      )
      filename = tmp_path / "checkpoint.pt"
      deepspeed_engine.save_checkpoint(filename, tag="checkpoint")
  
      model = ModelWithSharedWeights()
      state_dict = get_fp32_state_dict_from_zero_checkpoint(
          filename, tag="checkpoint"
      )
      model.load_state_dict(state_dict, strict=True)
  
file_path: DeepSpeed/tests/unit/checkpoint/test_shared_weights.py
imports: import torch
  import torch.nn
  import deepspeed
  from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
  from unit.common import DistributedTest
label_desc: Use image node_id 1 for calling the TestCheckpointSharedWeights obj's underlying member method code with example usage: obj.test_checkpoint_shared_weights(tmp_path) without return types