fix bug: _concatenate_results_from_multi_process

move all tensor to the same device
FlagOpen · Jan 13, 2025 · ddf9ada · ddf9ada
1 parent a572e63
commit ddf9ada
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/FlagEmbedding/abc/inference/AbsEmbedder.py b/FlagEmbedding/abc/inference/AbsEmbedder.py
@@ -416,6 +416,8 @@ def _concatenate_results_from_multi_process(self, results_list: List[Union[torch
             Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
         """
         if isinstance(results_list[0], torch.Tensor):
+            # move all tensors to the same device
+            results_list = [res.to(self.target_devices[0]) for res in results_list]
             return torch.cat(results_list, dim=0)
         elif isinstance(results_list[0], np.ndarray):
             return np.concatenate(results_list, axis=0)