diff --git a/paddle/fluid/inference/tensorrt/convert/multihead_matmul_op.cc b/paddle/fluid/inference/tensorrt/convert/multihead_matmul_op.cc
index 000a457a1a4536d0356a177faffd7422f413e5ce..8965be5fd8ec5a756f40c3016b136d051ca7e6fc 100644
--- a/paddle/fluid/inference/tensorrt/convert/multihead_matmul_op.cc
+++ b/paddle/fluid/inference/tensorrt/convert/multihead_matmul_op.cc
@@ -674,6 +674,8 @@ class MultiheadMatMulOpConverter : public OpConverter {
 
           // return
           layer = reshape_after_mha_layer;
+          RreplenishLayerAndOutput(
+              layer, "multihead_matmul", {output_name}, test_mode);
         } else {
           PADDLE_ENFORCE_EQ(
               input->getDimensions().nbDims,