code: def test_zero_fragments(
      self, tmpdir, api_type, zero_stage, offload_device, frozen_weights
  ):
      if offload_device == OffloadDeviceEnum.nvme:
          if zero_stage != 3:
              pytest.skip(
                  f"Nvme offload not supported for zero stage {zero_stage}"
              )
          if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
              pytest.skip("Skip tests since async-io is not compatible")
  
      if api_type == "local" and zero_stage != 3:
          pytest.skip(
              f"Local APIs only for zero stage 3 but current stage is {zero_stage}"
          )
  
      config_dict = {
          "train_micro_batch_size_per_gpu": 1,
          "steps_per_print": 1,
          "optimizer": {"type": "Adam", "params": {"lr": 1e-6}},
          "fp16": {"enabled": True, "initial_scale_power": 2},
          "zero_optimization": {
              "stage": zero_stage,
          },
      }
  
      if offload_device == OffloadDeviceEnum.cpu:
          config_dict["zero_optimization"]["offload_optimizer"] = {
              "device": offload_device
          }
      elif offload_device == OffloadDeviceEnum.nvme:
          config_dict["zero_optimization"]["offload_optimizer"] = {
              "device": offload_device,
              "nvme_path": str(tmpdir),
          }
  
      hidden_dim = 128
      if zero_stage == 3:
          with deepspeed.zero.Init(config_dict_or_path=config_dict):
              model = MyModel(hidden_dim, frozen_weights)
      else:
          model = MyModel(hidden_dim, frozen_weights)
  
      validate_func = validate_funcs_mapping[api_type]
  
      run_fragmented_model(
          model, config_dict, hidden_dim, torch.float16, validate_func
      )
  
imports: import pytest
  import deepspeed.comm
  import torch
  from unit.common import DistributedTest
  from unit.simple_model import random_dataloader, SimpleModel
  from unit.util import bf16_required_version_check
  import deepspeed
  from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state
  from deepspeed.utils import safe_set_full_fp32_param, safe_set_full_optimizer_state
  from deepspeed.utils import safe_get_local_fp32_param, safe_get_local_grad, safe_get_local_optimizer_state
  from deepspeed.utils import safe_set_local_fp32_param, safe_set_local_optimizer_state
  from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
  from deepspeed.ops.aio import AsyncIOBuilder
label_desc: Use image node_id 1 for calling the TestTensorFragmentGet obj's underlying member method code with example usage: obj.test_zero_fragments(tmpdir, api_type, zero_stage, offload_device, frozen_weights) without return types