def __call__(self, node_ids, device=th.device("cpu")):
        """
        node_ids : th.tensor
            Index of the embeddings to collect.
        device : th.device
            Target device to put the collected embeddings.
        """
        if not self._comm:
            # For embeddings stored on the CPU.
            emb = self._tensor[node_ids].to(device)
        else:
            # For embeddings stored on the GPU.
            # The following method is designed to perform communication
            # across multiple GPUs and can handle situations where only one GPU
            # is present gracefully, a.k.a. self._world_size == 1 or
            # 0 (when th.distributed.is_initialized() is false).
            emb = nccl.sparse_all_to_all_pull(
                node_ids, self._tensor, self._partition
            )
            emb = emb.to(device)
        if F.is_recording():
            emb = F.attach_grad(emb)
            self._trace.append((node_ids.to(device), emb))

        return emb

        node_ids : th.tensor
            Index of the embeddings to collect.
        device : th.device
            Target device to put the collected embeddings.
        
python/dgl/nn/pytorch/sparse_emb.py