changelog

Author: rusty1s

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [2.0.5] - 2022-MM-DD
 ### Added
+- Added a `filter_per_worker` argument to data loaders to allow filtering of data within sub-processes ([#4873](https://github.com/pyg-team/pytorch_geometric/pull/4873))
 - Added a `NeighborLoader` benchmark script ([#4815](https://github.com/pyg-team/pytorch_geometric/pull/4815))
 - Added support for `FeatureStore` and `GraphStore` in `NeighborLoader` ([#4817](https://github.com/pyg-team/pytorch_geometric/pull/4817), [#4851](https://github.com/pyg-team/pytorch_geometric/pull/4851), [#4854](https://github.com/pyg-team/pytorch_geometric/pull/4854), [#4856](https://github.com/pyg-team/pytorch_geometric/pull/4856), [#4857](https://github.com/pyg-team/pytorch_geometric/pull/4857))
 - Added a `normalize` parameter to `dense_diff_pool` ([#4847](https://github.com/pyg-team/pytorch_geometric/pull/4847))

torch_geometric/loader/hgt_loader.py

@@ -130,7 +130,7 @@ def __init__(
         self.colptr_dict, self.row_dict, self.perm_dict = to_hetero_csc(
             data, device='cpu', share_memory=kwargs.get('num_workers', 0) > 0)
 
-        super().__init__(input_nodes[1].tolist(), collate_fn=self.sample,
+        super().__init__(input_nodes[1].tolist(), collate_fn=self.collate_fn,
                          **kwargs)
 
     def sample(self, indices: List[int]) -> HeteroData:

torch_geometric/loader/link_neighbor_loader.py

@@ -266,7 +266,7 @@ def __init__(
         self.directed = directed
         self.neg_sampling_ratio = neg_sampling_ratio
         self.transform = transform
-        self.neighbor_sampler = neighbor_sampler
+        self.filter_per_worker = filter_per_worker
         self.neighbor_sampler = neighbor_sampler
 
         edge_type, edge_label_index = get_edge_label_index(