diff --git a/paddlenlp/transformers/tensor_parallel_utils.py b/paddlenlp/transformers/tensor_parallel_utils.py index 267c2846aa27..ea68e41207c5 100644 --- a/paddlenlp/transformers/tensor_parallel_utils.py +++ b/paddlenlp/transformers/tensor_parallel_utils.py @@ -44,7 +44,7 @@ def parallel_matmul(lm_output, logit_weights, tensor_parallel_output=True, train logits = xpu_parallel_matmul( lm_output, logit_weights, - tensor_parallel_output, + tensor_parallel_output=tensor_parallel_output, training=training, ) return logits