From 1a533ed2de02ce5541a3f7adaf3e00c1ffae3fe4 Mon Sep 17 00:00:00 2001
From: Zhaolong Xing <nhzlx.dragon@gmail.com>
Date: Thu, 5 Mar 2020 10:25:19 +0800
Subject: [PATCH] [BUG]: Multihead matmul op's ouput size should be BxSx(N*H)
 (#22848)

test=develop
---
 paddle/fluid/operators/fused/multihead_matmul_op.cu | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/paddle/fluid/operators/fused/multihead_matmul_op.cu b/paddle/fluid/operators/fused/multihead_matmul_op.cu
index 2500f66c67..de40ded24e 100644
--- a/paddle/fluid/operators/fused/multihead_matmul_op.cu
+++ b/paddle/fluid/operators/fused/multihead_matmul_op.cu
@@ -440,13 +440,11 @@ class MultiHeadMatMulV2Kernel : public framework::OpKernel<T> {
 
     auto &bias_qk = detail::Ref(context.Input<framework::Tensor>("BiasQK"),
                                 "Cannot find QK");
-    auto *out = context.Output<framework::Tensor>("Out");
 
     auto *input_d = input->data<T>();
     auto *w_d = w->data<T>();
     auto *bias_d = bias->data<T>();
     auto *bias_qk_d = bias_qk.data<T>();
-    auto *output_d = out->mutable_data<T>(context.GetPlace());
     T scale = static_cast<T>(context.Attr<float>("alpha"));
 
     int head_number = context.Attr<int>("head_number");
@@ -463,6 +461,10 @@ class MultiHeadMatMulV2Kernel : public framework::OpKernel<T> {
     int all_head_size = w_dims[2];
     int head_size = all_head_size / head_number;
 
+    auto *out = context.Output<framework::Tensor>("Out");
+    out->Resize({batch, seq_len, all_head_size});
+    auto *output_d = out->mutable_data<T>(context.GetPlace());
+
     // (B*S, hidden)
     const Tensor input_matrix =
         framework::ReshapeToMatrix(*input, 2 /*x_num_col_dims */);
-- 
GitLab