diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.cc b/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.cc
index 9d6a973323ef4307a452d7d94dab71541a52d6da..74be7c143ef8f3d5cf33fa2bc0eb5c000691aae7 100644
--- a/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.cc
+++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.cc
@@ -165,7 +165,7 @@ std::vector<DistTensorSpec> MatmulSPMDRule::InferForward(
 
 TensorDistAttr GetInferedDistAttr(
     const TensorDistAttr& origin_dist_attr,
-    const std::vector<int>& shape,
+    const std::vector<int64_t>& shape,
     const std::string& tensor_axis,
     const std::unordered_map<std::string, int64_t>& axis_to_dim_map) {
   TensorDistAttr dist_attr_ = CopyTensorDistAttrForOutput(origin_dist_attr);
diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.h b/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.h
index 89677e8e4e291a1a3433d4403ae59fa62ca91726..65b0600eae5c24731c197a745ffcc5fbf45dfdae 100644
--- a/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.h
+++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.h
@@ -27,7 +27,7 @@ namespace auto_parallel {
 
 TensorDistAttr GetInferedDistAttr(
     const TensorDistAttr& origin_dist_attr,
-    const std::vector<int>& shape,
+    const std::vector<int64_t>& shape,
     const std::string& tensor_axis,
     const std::unordered_map<std::string, int64_t>& axis_to_dim_map);