From 275183b16bb2333bb82774c2273b69c12292ac9f Mon Sep 17 00:00:00 2001 From: Andrei Ivanov <32910461+drivanov@users.noreply.github.com> Date: Sun, 5 Jan 2025 18:34:28 -0800 Subject: [PATCH] [Test] Fixing issues related to `SampledHeteroCSC` and `StaticHeteroCSC` deprecation. (#7799) --- python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py b/python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py index 3b1adc7309..41114e1889 100644 --- a/python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py +++ b/python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py @@ -9,7 +9,7 @@ from torch import nn from .cugraph_base import CuGraphBaseConv try: - from pylibcugraphops.pytorch import SampledHeteroCSC, StaticHeteroCSC + from pylibcugraphops.pytorch import HeteroCSC from pylibcugraphops.pytorch.operators import ( agg_hg_basis_n2n_post as RelGraphConvAgg, ) @@ -188,27 +188,28 @@ class CuGraphRelGraphConv(CuGraphBaseConv): max_in_degree = g.in_degrees().max().item() if max_in_degree < self.MAX_IN_DEGREE_MFG: - _graph = SampledHeteroCSC( + _graph = HeteroCSC( offsets, indices, edge_types_perm, - max_in_degree, g.num_src_nodes(), self.num_rels, ) else: offsets_fg = self.pad_offsets(offsets, g.num_src_nodes() + 1) - _graph = StaticHeteroCSC( + _graph = HeteroCSC( offsets_fg, indices, edge_types_perm, + g.num_src_nodes(), self.num_rels, ) else: - _graph = StaticHeteroCSC( + _graph = HeteroCSC( offsets, indices, edge_types_perm, + g.num_src_nodes(), self.num_rels, )