Method multi_gpu_encode

src/flashrag/retriever/encoder.py:79–83

(self, query_list: Union[List[str], str], batch_size=64, is_query=True)

    @torch.inference_mode()
    def multi_gpu_encode(self, query_list: Union[List[str], str], batch_size=64, is_query=True) -> np.ndarray:
        if self.gpu_num > 1:
            self.model = torch.nn.DataParallel(self.model)
        query_emb = self.encode(query_list, batch_size, is_query)
        return query_emb


class STEncoder: