def _all_set_optm_state(self, states):
        """Set the optimizer states of the embedding. This method must be
        called by all processes sharing the embedding with identical
        :attr:`states`.

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

        Parameters
        ----------
        states : tuple of torch.Tensor
            The global states to pull values from.
        """
        if self._partition:
            idxs = F.copy_to(
                self._partition.get_local_indices(
                    max(self._rank, 0), ctx=F.context(self._tensor)
                ),
                F.context(states[0]),
            )
            for state, new_state in zip(self._optm_state, states):
                state[:] = F.copy_to(
                    F.gather_row(new_state, idxs), ctx=F.context(self._tensor)
                )[:]
        else:
            # stored in CPU memory
            if self._rank <= 0:
                for state, new_state in zip(self._optm_state, states):
                    state[:] = F.copy_to(
                        new_state, ctx=F.context(self._tensor)
                    )[:]
        if th.distributed.is_initialized():
            th.distributed.barrier()
Set the optimizer states of the embedding. This method must be
        called by all processes sharing the embedding with identical
        :attr:`states`.

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

        Parameters
        ----------
        states : tuple of torch.Tensor
            The global states to pull values from.
        
python/dgl/nn/pytorch/sparse_emb.py