From 8772c02db74500096174e87f41a7f47ff2fa2248 Mon Sep 17 00:00:00 2001
From: Rhett Ying <85214957+Rhett-Ying@users.noreply.github.com>
Date: Fri, 6 Sep 2024 13:56:20 +0800
Subject: [PATCH] [DistGB] fix prob attr not exist (#7783)

---
 python/dgl/distributed/dist_dataloader.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/python/dgl/distributed/dist_dataloader.py b/python/dgl/distributed/dist_dataloader.py
index d2ce3f52177d..117379b2b4a3 100644
--- a/python/dgl/distributed/dist_dataloader.py
+++ b/python/dgl/distributed/dist_dataloader.py
@@ -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):
@@ -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):