Fix pyg-lib hetero_neighbor_sampler usage (#5697)

Co-authored-by: Matthias Fey <[email protected]>
pyg-team · Oct 15, 2022 · 753d1e6 · 753d1e6
1 parent a4ff7a7
commit 753d1e6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_geometric/sampler/neighbor_sampler.py b/torch_geometric/sampler/neighbor_sampler.py
@@ -227,7 +227,7 @@ def _sample(
                 # TODO (matthias) Add `disjoint` option to `NeighborSampler`
                 # TODO (matthias) `return_edge_id` if edge features present
                 disjoint = self.node_time_dict is not None
-                out = torch.ops.pyg.hetero_neighbor_sample_cpu(
+                out = torch.ops.pyg.hetero_neighbor_sample(
                     self.node_types,
                     self.edge_types,
                     self.colptr_dict,