Skip to content

Commit

Permalink
[DistGB] fix prob attr not exist (#7783)
Browse files Browse the repository at this point in the history
  • Loading branch information
Rhett-Ying authored Sep 6, 2024
1 parent f9097ee commit 8772c02
Showing 1 changed file with 8 additions and 2 deletions.
10 changes: 8 additions & 2 deletions python/dgl/distributed/dist_dataloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -378,7 +378,10 @@ def __init__(self, g, nids, graph_sampler):
self._dataset = utils.maybe_flatten_dict(self.nids)

# Add prob/mask into graphbolt partition's edge attributes if needed.
Collator.add_edge_attribute_to_graph(self.g, self.graph_sampler.prob)
if hasattr(self.graph_sampler, "prob"):
Collator.add_edge_attribute_to_graph(
self.g, self.graph_sampler.prob
)

@property
def dataset(self):
Expand Down Expand Up @@ -637,7 +640,10 @@ def __init__(
self._dataset = utils.maybe_flatten_dict(self.eids)

# Add prob/mask into graphbolt partition's edge attributes if needed.
Collator.add_edge_attribute_to_graph(self.g, self.graph_sampler.prob)
if hasattr(self.graph_sampler, "prob"):
Collator.add_edge_attribute_to_graph(
self.g, self.graph_sampler.prob
)

@property
def dataset(self):
Expand Down

0 comments on commit 8772c02

Please sign in to comment.