Skip to content

Commit

Permalink
[Test] Fixing issues related to SampledHeteroCSC and `StaticHetero…
Browse files Browse the repository at this point in the history
…CSC` deprecation. (#7799)
  • Loading branch information
drivanov authored Jan 6, 2025
1 parent 88f109f commit 275183b
Showing 1 changed file with 6 additions and 5 deletions.
11 changes: 6 additions & 5 deletions python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
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,
)
Expand Down Expand Up @@ -188,27 +188,28 @@ def forward(self, g, feat, etypes, max_in_degree=None):
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,
)

Expand Down

0 comments on commit 275183b

Please sign in to comment.