Fix index-out-of-bounds issue in LinkPredNDCG(weighted=True) (#9963)

pyg-team · Jan 21, 2025 · 47ac818 · 47ac818
1 parent 9c73f99
commit 47ac818
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_geometric/metrics/link_pred.py b/torch_geometric/metrics/link_pred.py
@@ -492,7 +492,7 @@ def _compute(self, data: LinkPredMetricData) -> Tensor:
                 self.discount,
                 self.discount.new_full((1, ), fill_value=float('inf')),
             ])
-            discount = discount[pos.clamp(max=self.k + 1)]
+            discount = discount[pos.clamp(max=self.k)]
 
             idcg = scatter(  # Apply discount and aggregate:
                 data.edge_label_weight / discount,