Unverified Commit cf349c4a authored by Isotr0py's avatar Isotr0py Committed by GitHub
Browse files

[Bugfix][CPU] Fix CPU embedding runner with tensor parallel (#10394)


Signed-off-by: default avatarIsotr0py <2037008807@qq.com>
parent 905d0f0a
...@@ -66,6 +66,10 @@ class CPUEmbeddingModelRunner( ...@@ -66,6 +66,10 @@ class CPUEmbeddingModelRunner(
hidden_states = model_executable(**execute_model_kwargs) hidden_states = model_executable(**execute_model_kwargs)
# Only perform pooling in the driver worker.
if not self.is_driver_worker:
return []
return [ return [
self.model.pooler(hidden_states=hidden_states, self.model.pooler(hidden_states=hidden_states,
pooling_metadata=model_input.pooling_metadata) pooling_metadata=model_input.pooling_metadata)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment