Minor fix

deepset-ai · Oct 7, 2023 · 6d7f7e1 · 6d7f7e1
1 parent 5a6f341
commit 6d7f7e1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/haystack/preview/components/samplers/top_p.py b/haystack/preview/components/samplers/top_p.py
@@ -104,7 +104,7 @@ def run(self, documents: List[Document], top_p: Optional[float] = None):
                 top_p,
             )
             highest_prob_indices = torch.argsort(probs, descending=True)
-            selected_docs = [documents[int(highest_prob_indices[0].item())]]
+            selected_docs = [documents[highest_prob_indices[0].item()]]
 
         return {"documents": selected_docs}