Fix homo temporal distributed sampling (#8367)

First get node attributes and then extract node_types to prevent node_attrs from being always be empty in case of homo sampling. --------- Co-authored-by: rusty1s <[email protected]>
pyg-team · Nov 13, 2023 · b1f8535 · b1f8535
1 parent 9f7e824
commit b1f8535
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -23,7 +23,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
-- Update `DistNeighborSampler` for homogeneous graphs ([#8209](https://github.com/pyg-team/pytorch_geometric/pull/8209))
+- Update `DistNeighborSampler` for homogeneous graphs ([#8209](https://github.com/pyg-team/pytorch_geometric/pull/8209), [#8367](https://github.com/pyg-team/pytorch_geometric/pull/8367))
 - Update `GraphStore` and `FeatureStore` to support distributed training ([#8083](https://github.com/pyg-team/pytorch_geometric/pull/8083))
 - Disallow the usage of `add_self_loops=True` in `GCNConv(normalize=False)` ([#8210](https://github.com/pyg-team/pytorch_geometric/pull/8210))
 - Disable device asserts during `torch_geometric.compile` ([#8220](https://github.com/pyg-team/pytorch_geometric/pull/8220))

diff --git a/torch_geometric/sampler/neighbor_sampler.py b/torch_geometric/sampler/neighbor_sampler.py
@@ -121,7 +121,8 @@ def __init__(
             # Obtain graph metadata:
             node_attrs = [
                 attr for attr in feature_store.get_all_tensor_attrs()
-                if isinstance(attr.group_name, NodeType)
+                if isinstance(attr.group_name, NodeType)  # Heterogeneous ...
+                or attr.group_name is None  # ... or homogeneous.
             ]
             self.node_types = list(set(attr.group_name for attr in node_attrs))