Skip to content

Commit

Permalink
Merge pull request #1 from eltociear/patch-1
Browse files Browse the repository at this point in the history
chore: update auto_model.py
  • Loading branch information
archersama authored Jul 5, 2024
2 parents 931b066 + d67fe7f commit 5a91eaa
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions coir/beir/generation/models/auto_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ def start_multi_process_pool(self, target_devices: List[str] = None):
This method is recommended if you want to encode on multiple GPUs. It is advised
to start only one process per GPU. This method works together with encode_multi_process
:param target_devices: PyTorch target devices, e.g. cuda:0, cuda:1... If None, all available CUDA devices will be used
:return: Returns a dict with the target processes, an input queue and and output queue.
:return: Returns a dict with the target processes, an input queue and output queue.
"""
if target_devices is None:
if torch.cuda.is_available():
Expand Down Expand Up @@ -158,4 +158,4 @@ def generate_multi_process(self, corpus: List[Dict[str, str]], ques_per_passage:
results_list = sorted([output_queue.get() for _ in range(last_chunk_id)], key=lambda x: x[0])
queries = [result[1] for result in results_list]

return [item for sublist in queries for item in sublist]
return [item for sublist in queries for item in sublist]

0 comments on commit 5a91eaa

Please sign in to comment.