code: def test(self):
      config_dict = {
          "train_batch_size": 2,
          "steps_per_print": 1,
          "sparse_gradients": True,
      }
  
      model = Model()
      optimizer = Adam(
          list(model.linear.parameters()), list(model.emb.parameters())
      )
      engine, _, _, _ = deepspeed.initialize(
          model=model, optimizer=optimizer, config=config_dict
      )
      loss = torch.nn.BCEWithLogitsLoss()
      x = torch.tensor(
          [1, 2, 4, 5, 4, 3, 2, 9],
          dtype=torch.long,
          device=engine.device,
      )
      offsets = torch.tensor(
          [0, 4], dtype=torch.long, device=engine.device
      )
      y = torch.tensor([[1.0], [0.0]], device=engine.device)
      res = engine(x, offsets)
      engine.backward(loss(res, y))
      engine.step()
  
      results = [
          engine.all_gather_scalar(i, groups._get_data_parallel_group())
          for i in model.emb.parameters()
      ]
      for res in results:
          assert torch.allclose(res[0], res[1])
  
file_path: deepspeed/tests/unit/runtime/sparse_tensor/test_sparse_grads.py
imports: import torch
  import deepspeed
  from unit.common import DistributedTest
  import deepspeed.utils.groups
label_desc: Use image node_id 1 for calling the TestSparseAdam obj's underlying member method code with example usage: obj.test() without return types