From 6d7f7e1dbad80100692cf5a43ffbb36faae635f8 Mon Sep 17 00:00:00 2001 From: Vladimir Blagojevic Date: Sat, 7 Oct 2023 11:33:11 +0200 Subject: [PATCH] Minor fix --- haystack/preview/components/samplers/top_p.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/haystack/preview/components/samplers/top_p.py b/haystack/preview/components/samplers/top_p.py index 6f3fe57772..bb5eb27082 100644 --- 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}