Explicitly add reserve_space output tensor

PaddlePaddle · Nov 14, 2023 · fccd5ee · fccd5ee
1 parent 981ca69
commit fccd5ee
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/paddle/fluid/pir/transforms/fusion/fused_gemm_epilogue_pass.cc b/paddle/fluid/pir/transforms/fusion/fused_gemm_epilogue_pass.cc
@@ -143,7 +143,7 @@ class FusedLinearGeluPattern
                  {"activation", act_attr}}});
     fused_gemm_epilogue_gelu(
         {&res.Tensor("x"), &res.Tensor("w"), &res.Tensor("bias")},
-        {&res.Tensor("out")});
+        {&res.Tensor("out"), &res.Tensor("reserve_space")});
   }
 };
 class FusedLinearReluPattern
@@ -181,7 +181,7 @@ class FusedLinearReluPattern
                  {"activation", act_attr}}});
     fused_gemm_epilogue_relu(
         {&res.Tensor("x"), &res.Tensor("w"), &res.Tensor("bias")},
-        {&res.Tensor("out")});
+        {&res.Tensor("out"), &res.Tensor("reserve_space")});
   }
 };