class_docstr: The abstract sparse optimizer.
  
  Note: dgl sparse optimizer only work with dgl.NodeEmbedding
  
  Parameters
  ----------
  params : list of NodeEmbedding
      The list of NodeEmbeddings.
  lr : float
      The learning rate.
docstr: The step function.
  
  The step function is invoked at the end of every batch to update embeddings
docstr_tok: ["The","step","function",".","The","step","function","is","invoked","at","the","end","of","every","batch","to","update","embeddings"]
code: def step(self):
      """The step function.
  
      The step function is invoked at the end of every batch to update embeddings
      """
      # on the first step, check to see if the grads are on the GPU
      if self._first_step:
          for emb in self._params:
              for _, data in emb._trace:
                  if data.grad.device.type == "cuda":
                      # create a communicator
                      if self._device:
                          assert (
                              self._device == data.grad.device
                          ), "All gradients must be on the same device"
                      else:
                          self._device = data.grad.device
                  else:
                      assert (
                          not self._device
                      ), "All gradients must be on the same device"
  
          # distributed backend use nccl
          if self._device and (
              not th.distributed.is_initialized()
              or th.distributed.get_backend() == "nccl"
          ):
              # device is only set if the grads are on a GPU
              self._comm_setup()
          else:
              self._shared_setup()
          self._first_step = False
  
      if self._comm:
          self._comm_step()
      else:
          self._shared_step()
  
imports: import abc
  from abc import abstractmethod
  import torch
  from ...cuda import nccl
  from ...nn.pytorch import NodeEmbedding
  from ...partition import NDArrayPartition
  from ...utils import create_shared_mem_array, gather_pinned_tensor_rows, get_shared_mem_array, pin_memory_inplace, scatter_pinned_tensor_rows
label_desc: Use image node_id 2 for calling the SparseGradOptimizer obj's underlying member method code with example usage: obj.step() without return types