fix miss input param for _fuse_prepare_qkv (PaddlePaddle#655)

heavengate · Aug 23, 2022 · 865abf4 · 865abf4
1 parent 0b73295
commit 865abf4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/fleetx/models/gpt_model/modeling_auto.py b/fleetx/models/gpt_model/modeling_auto.py
@@ -80,7 +80,7 @@ def __init__(self,
         self.out_proj = nn.Linear(
             embed_dim, embed_dim, weight_attr, bias_attr=bias_attr)
 
-    def _fuse_prepare_qkv(self, query):
+    def _fuse_prepare_qkv(self, query, use_cache=False, cache=None):
         mix_layer = self.qkv_proj(query)
         mix_layer = paddle.reshape_(mix_layer,
                                     [0, 0, self.num_heads, 3 * self.head_dim])

diff --git a/fleetx/models/gpt_model/modeling_hybrid.py b/fleetx/models/gpt_model/modeling_hybrid.py
@@ -151,7 +151,7 @@ def __init__(self,
             input_is_parallel=True,
             fuse_matmul_bias=fused_linear)
 
-    def _fuse_prepare_qkv(self, query):
+    def _fuse_prepare_qkv(self, query, use_cache=False, cache=None):
         mix_layer = self.qkv_proj(query)
         mix_layer = paddle.reshape_(mix_layer,
                                     [0, 0, self.num_heads, 3 * self.head_dim])