def _all_get_optm_state(self):
        """Return a copy of the whole optimizer states stored in CPU memory.
        If this is a multi-processing instance, the states will be returned in
        shared memory. If the embedding is currently stored on multiple GPUs,
        all processes must call this method in the same order.

        NOTE: This method must be called by all processes sharing the
        embedding, or it may result in a deadlock.

        Returns
        -------
        tuple of torch.Tensor
            The optimizer states stored in CPU memory.
        """
        if self._partition:
            if self._world_size == 0:
                # non-multiprocessing
                return tuple(
                    state.to(th.device("cpu")) for state in self._optm_state
                )
            else:
                return tuple(
                    self._all_get_tensor(
                        f"state_gather_{self._name}_{i}",
                        state,
                        (self._num_embeddings, *state.shape[1:]),
                    )
                    for i, state in enumerate(self._optm_state)
                )
        else:
            # already stored in CPU memory
            return self._optm_state
Return a copy of the whole optimizer states stored in CPU memory.
        If this is a multi-processing instance, the states will be returned in
        shared memory. If the embedding is currently stored on multiple GPUs,
        all processes must call this method in the same order.

        NOTE: This method must be called by all processes sharing the
        embedding, or it may result in a deadlock.

        Returns
        -------
        tuple of torch.Tensor
            The optimizer states stored in CPU memory.
        
python/dgl/nn/pytorch/sparse_emb.py