[Bug Fix] Update dist_neighbor_sampler to handle heterogeneous case where node labels is None (#153)

mkolodner-sc · web-flow · commit cb61c2734cf4 · 2025-04-01T10:58:37.000+08:00
diff --git a/graphlearn_torch/python/distributed/dist_neighbor_sampler.py b/graphlearn_torch/python/distributed/dist_neighbor_sampler.py
@@ -182,7 +182,7 @@ def __init__(self,
             local_only=False, rpc_router=self.rpc_router, device=self.device
           )
       else:
-        assert isinstance(self.dist_node_labels, Dict)
+        assert self.dist_node_labels is None or isinstance(self.dist_node_labels, Dict)
         if self.dist_node_labels is not None and \
             all(isinstance(value, Feature) for value in self.dist_node_labels.values()):
           self.dist_node_labels = DistFeature(

Original file line number	Diff line number	Diff line change
`@@ -182,7 +182,7 @@ def __init__(self,`
`182`	`182`	`local_only=False, rpc_router=self.rpc_router, device=self.device`
`183`	`183`	`)`
`184`	`184`	`else:`
`185`		`- assert isinstance(self.dist_node_labels, Dict)`
	`185`	`+ assert self.dist_node_labels is None or isinstance(self.dist_node_labels, Dict)`
`186`	`186`	`if self.dist_node_labels is not None and \`
`187`	`187`	`all(isinstance(value, Feature) for value in self.dist_node_labels.values()):`
`188`	`188`	`self.dist_node_labels = DistFeature(`