patch

alexbarghi-nv · Aug 1, 2024 · 34a1ecd · 34a1ecd
1 parent f4c519e
commit 34a1ecd
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
@@ -157,7 +157,7 @@ def __init__(self, memory_type="distributed", location="cpu"):
 
         self.__features = {}
 
-        self.__wg_comm = wgth.get_local_node_communicator()
+        self.__wg_comm = wgth.get_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
@@ -260,7 +260,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,
             )
         )