diff --git a/paddle/fluid/operators/ngraph/ops/softmax_op.h b/paddle/fluid/operators/ngraph/ops/softmax_op.h
index 174b7a91a8dd0e3edb06f224c3914e24c6c4a96d..6eb84703998c24ee7b9e0d4f6931c3fe0bd00e2e 100644
--- a/paddle/fluid/operators/ngraph/ops/softmax_op.h
+++ b/paddle/fluid/operators/ngraph/ops/softmax_op.h
@@ -35,7 +35,7 @@ std::shared_ptr<ngraph::Node> GetSoftmax(std::shared_ptr<ngraph::Node> x) {
 
   auto x_max = std::make_shared<ngraph::op::Max>(x, ngraph::AxisSet{1});
   auto x_max_bcast = std::make_shared<ngraph::op::Broadcast>(
-      x_max, x_shape, ngraph::AxisSet{1});
+      x_max, x_2d_shape, ngraph::AxisSet{1});
   auto x_shifted = x - x_max_bcast;
   auto x_clipped =
       paddle::operators::ngraphs::ElementwiseScalar<ngraph::op::Maximum>(