diff --git a/paddle/framework/ddim.cc b/paddle/framework/ddim.cc
index 9431645cf5b290ee7638bc1970c4cc76891a2aa0..3fd3e538e8da41bb6fbf533dad3789688a734e0d 100644
--- a/paddle/framework/ddim.cc
+++ b/paddle/framework/ddim.cc
@@ -226,7 +226,7 @@ Eigen::DSizes<Eigen::DenseIndex, NDIMS> ToEigenDSizes(DDim dims) const {
   int rank = paddle::framework::arity(dims);
   PADDLE_ENFORCE(rank == NDIMS, "DDim and NDIMS must be same")
   Eigen::DSizes<Eigen::DenseIndex, NDIMS> dsizes;
-  for (int d = 0; d < paddle::framework::arity(dims); d++) {
+  for (int d = 0; d < rank; d++) {
     dsizes[d] = dims[d];
   }
   return dsizes;