diff --git a/test/utils/test_convert.py b/test/utils/test_convert.py index 93ea55d84586..8e4b8280a98f 100644 --- a/test/utils/test_convert.py +++ b/test/utils/test_convert.py @@ -462,13 +462,13 @@ def test_to_cugraph(edge_weight, directed, relabel_nodes): edge_list = graph.view_edge_list() assert edge_list is not None - edge_list = edge_list.sort_values(by=['src', 'dst']) + edge_list = edge_list.sort_values(by=[0, 1]) - cu_edge_index = edge_list[['src', 'dst']].to_pandas().values + cu_edge_index = edge_list[[0, 1]].to_pandas().values cu_edge_index = torch.from_numpy(cu_edge_index).t() cu_edge_weight = None if edge_weight is not None: - cu_edge_weight = edge_list['weights'].to_pandas().values + cu_edge_weight = edge_list['2'].to_pandas().values cu_edge_weight = torch.from_numpy(cu_edge_weight) cu_edge_index, cu_edge_weight = sort_edge_index(cu_edge_index, diff --git a/torch_geometric/utils/convert.py b/torch_geometric/utils/convert.py index eccb23051bbb..e0fa499b5799 100644 --- a/torch_geometric/utils/convert.py +++ b/torch_geometric/utils/convert.py @@ -457,13 +457,13 @@ def from_cugraph(g: Any) -> Tuple[Tensor, Optional[Tensor]]: """ df = g.view_edge_list() - src = from_dlpack(df['src'].to_dlpack()).long() - dst = from_dlpack(df['dst'].to_dlpack()).long() + src = from_dlpack(df[0].to_dlpack()).long() + dst = from_dlpack(df[1].to_dlpack()).long() edge_index = torch.stack([src, dst], dim=0) edge_weight = None - if 'weights' in df: - edge_weight = from_dlpack(df['weights'].to_dlpack()) + if '2' in df: + edge_weight = from_dlpack(df['2'].to_dlpack()) return edge_index, edge_weight