From d0751d09fa61d5b4798b927e52f27e03b2f34b8b Mon Sep 17 00:00:00 2001
From: Shang Zhizhou <shangzhizhou@baidu.com>
Date: Mon, 26 Apr 2021 09:43:32 +0800
Subject: [PATCH] fix ernie oss error (#32540)

---
 paddle/fluid/inference/tensorrt/convert/fc_op.cc | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/paddle/fluid/inference/tensorrt/convert/fc_op.cc b/paddle/fluid/inference/tensorrt/convert/fc_op.cc
index aebdb8f884c..6167e68df2b 100644
--- a/paddle/fluid/inference/tensorrt/convert/fc_op.cc
+++ b/paddle/fluid/inference/tensorrt/convert/fc_op.cc
@@ -160,6 +160,12 @@ class FcOpConverter : public OpConverter {
     if (engine_->with_dynamic_shape()) {
       // not NCHW layout, but NLP layout with added 'x 1 x 1'
       auto x_dim = X->getDimensions();
+      if (engine_->use_oss() && engine_->with_ernie() && x_dim.nbDims == 4 &&
+          x_dim.d[2] == 1 && x_dim.d[3] == 1 && x_num_col_dims == 2) {
+        // fc which is just after self attention
+        regist_fc(X, n_output, weight, bias);
+        return;
+      }
       PADDLE_ENFORCE_LE(
           x_dim.nbDims - x_num_col_dims, 3,
           platform::errors::InvalidArgument(
-- 
GitLab