diff --git a/torch_geometric/utils/rag/feature_store.py b/torch_geometric/utils/rag/feature_store.py index e01e9e59bb88..f36b885d9fea 100644 --- a/torch_geometric/utils/rag/feature_store.py +++ b/torch_geometric/utils/rag/feature_store.py @@ -86,8 +86,9 @@ def load_subgraph( # NOTE: torch_geometric.loader.utils.filter_custom_store can be used # here if it supported edge features node_id = sample.node + print("sampled node id in load subgraph=", node_id) edge_id = sample.edge - edge_index = torch.stack((sample.row, sample.col), dim=0) + edge_index = torch.stack((sample.orig_row, sample.orig_col), dim=0) x = self.x[node_id] edge_attr = self.edge_attr[edge_id]