fp32 topk on cpu (#34)

* fp32 topk on cpu * black
graphcore-research · Oct 27, 2023 · 14f39f8 · 14f39f8
1 parent 0fdb741
commit 14f39f8
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/besskge/pipeline.py b/besskge/pipeline.py
@@ -271,7 +271,11 @@ def forward(self) -> Dict[str, Any]:
             if self.return_scores:
                 scores.append(batch_scores_filt)
             if self.return_topk:
-                topk_ids.append(torch.topk(batch_scores_filt, k=self.k, dim=-1).indices)
+                topk_ids.append(
+                    torch.topk(
+                        batch_scores_filt.to(torch.float32), k=self.k, dim=-1
+                    ).indices
+                )
 
         out = dict()
         if scores: