retriever.retrieve(corpus, queries) with len(queries)==1 errors

[manestay]

I have a somewhat silly use case. I'm running retriever.retrieve, with a queries dict with only 1 entry. However, this causes an IndexError with pytorch due to how pytorch indexes 2D arrays where the first dim is of size 1:

File "<dir>/script.py", line 83, in <module>
    results = retriever.retrieve(para_d, one_query_d)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<env>/beir/retrieval/evaluation.py", line 20, in retrieve
    return self.retriever.search(corpus, queries, self.top_k, self.score_function, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<env>/beir/retrieval/search/dense/exact_search.py", line 73, in search
    scores_dim = cos_scores[1]
                 ~~~~~~~~~~^^^
IndexError: index 1 is out of bounds for dimension 0 with size 1

I was able to fix this by changing the below line:

beir/beir/retrieval/search/dense/exact_search.py

Line 73 in f062f03

cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(cos_scores, min(top_k+1, len(cos_scores[1])), dim=1, largest=True, sorted=return_sorted)

And swapping out cos_scores[1] with scores_dim:

scores_dim = cos_scores[1] if cos_scores.shape[0] != 1 else cos_scores[0]

This monkey patch works for me, but I'm just curious if there's a more appropriate way to retrieve given only 1 query. Thanks!