"git@developer.sourcefind.cn:OpenDAS/vision.git" did not exist on "a89da92ba61355de1c9861338cd1d02ad5e97fff"
Unverified Commit a454734f authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] Follow up on the missing comment in PR#4989. (#5437)


Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>
Co-authored-by: default avatarXin Yao <yaox12@outlook.com>
parent 8d5d8962
...@@ -168,11 +168,14 @@ class NodeEmbedding: # NodeEmbedding ...@@ -168,11 +168,14 @@ class NodeEmbedding: # NodeEmbedding
Target device to put the collected embeddings. Target device to put the collected embeddings.
""" """
if not self._comm: if not self._comm:
# embeddings are stored on the CPU # For embeddings stored on the CPU.
emb = self._tensor[node_ids].to(device) emb = self._tensor[node_ids].to(device)
else: else:
# embeddings are stored on the GPU # For embeddings stored on the GPU.
# the following method also covers self._world_size = 0 or 1 # 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( emb = nccl.sparse_all_to_all_pull(
node_ids, self._tensor, self._partition node_ids, self._tensor, self._partition
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment