Merge pull request #20 from alexbarghi-nv/correct-wg-comm

[BUG] Use the Correct WG Communicator (rapidsai/cugraph#4548)
rapidsai · Aug 1, 2024 · 961fd04 · 961fd04
2 parents fca4b79 + 90db89a
commit 961fd04
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/python/cugraph-pyg/cugraph_pyg/data/feature_store.py b/python/cugraph-pyg/cugraph_pyg/data/feature_store.py
@@ -169,7 +169,7 @@ def __init__(self, memory_type="distributed", location="cpu"):
 
         self.__features = {}
 
-        self.__wg_comm = wgth.get_local_node_communicator()
+        self.__wg_comm = wgth.get_global_communicator()
         self.__wg_type = memory_type
         self.__wg_location = location
 

diff --git a/python/cugraph-pyg/cugraph_pyg/data/graph_store.py b/python/cugraph-pyg/cugraph_pyg/data/graph_store.py
@@ -271,7 +271,7 @@ def __get_edgelist(self):
             torch.tensor(
                 [self.__edge_indices[et].shape[1] for et in sorted_keys],
                 device="cuda",
-                dtype=torch.int32,
+                dtype=torch.int64,
             )
         )