diff --git a/paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc
index ef5d95dca3f63ca57725fe4d7a8ef3e23113fa1d..eb5059feaec05a6b64cf48aacc223db93dd05062 100644
--- a/paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc
+++ b/paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc
@@ -47,8 +47,12 @@ class SoftmaxMKLDNNHandler
         platform::errors::InvalidArgument(
             "The shape of input and output tensor must be identical."));
 
-    this->AcquireForwardPrimitiveDescriptor(prop_kind::forward_scoring,
-                                            input->mem_desc(), axis);
+    auto softmax_tz = phi::vectorize(input->dims());
+    auto md = memory::desc(softmax_tz, platform::MKLDNNGetDataType<T>(),
+                           input->format());
+
+    this->AcquireForwardPrimitiveDescriptor(prop_kind::forward_scoring, md,
+                                            axis);
   }
 
   SoftmaxMKLDNNHandler(const framework::ExecutionContext& ctx,