skip useless repeat_interleave computation in some specific but commo…

…n scenarios
pyg-team · Jan 7, 2025 · 06088d5 · 06088d5
1 parent bd296a3
commit 06088d5
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/torch_geometric/data/batch.py b/torch_geometric/data/batch.py
@@ -243,8 +243,9 @@ def filter(self, idx: torch.Tensor) -> Self:
             for attr, slc in slices.items():
                 slice_diff = slc.diff()
 
-                # Reshape mask to align it with attribute shape
-                attr_mask = mask[torch.repeat_interleave(slice_diff)]
+                # Reshape mask to align it with attribute shape.
+                # Since slice_diff often contains only ones, skip useless computation in such cases
+                attr_mask = mask[torch.repeat_interleave(slice_diff)] if torch.any(slice_diff != 1) else mask
 
                 # Apply mask to attribute
                 if attr == 'edge_index':