notation consistency

graphcore-research · Oct 24, 2023 · 282a0ea · 282a0ea
1 parent 4c3a92e
commit 282a0ea
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/besskge/pipeline.py b/besskge/pipeline.py
@@ -188,7 +188,7 @@ def forward(self) -> Dict[str, Any]:
 
         out = dict(scores=torch.concat(scores, dim=0))
         if ids:
-            out["triple_ids"] = torch.concat(ids, dim=0)
+            out["triple_idx"] = torch.concat(ids, dim=0)
         if self.evaluation:
             concat_metrics = dict()
             for m in metrics[0].keys():

diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
@@ -90,7 +90,7 @@ def test_all_scores_pipeline(corruption_scheme: str, compute_metrics: bool) -> N
 
     # Shuffle triples in same order of out["scores"]
     triple_reordered = torch.from_numpy(
-        ds.triples["test"][partitioned_triple_set.triple_sort_idx[out["triple_ids"]]]
+        ds.triples["test"][partitioned_triple_set.triple_sort_idx[out["triple_idx"]]]
     )
 
     # Real scores