fix op pd_op.mean

PaddlePaddle · Feb 7, 2025 · d111ea4 · d111ea4
1 parent 1c8c3fb
commit d111ea4
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/paddle2onnx/mapper/tensor/reduce_mean.cc b/paddle2onnx/mapper/tensor/reduce_mean.cc
@@ -38,7 +38,11 @@ void ReduceMeanMapper::Opset18() {
       GetAttr(axis_name_, &dim_);
     }
   } else {
-    GetAttr("axis", &dim_);
+    if (HasAttr(axis_name_)) {
+      GetAttr(axis_name_, &dim_);
+    } else {
+      TryGetInputValue(axis_name_, &dim_);
+    }
     if (dim_.size() == 0) {
       reduce_all_ = true;
     } else {
@@ -51,6 +55,9 @@ void ReduceMeanMapper::Opset18() {
   if (IsAttrVar(axis_name_)) {
     auto info = GetAttrVar(axis_name_);
     dims = helper_->AutoCast(info[0].name, info[0].dtype, P2ODataType::INT64);
+  } else if(HasInput(axis_name_)) {
+    auto info = GetInput(axis_name_);
+    dims = helper_->AutoCast(info[0].name, info[0].dtype, P2ODataType::INT64);
   } else {
     if (!reduce_all_) {
       dims = helper_->Constant(ONNX_NAMESPACE::TensorProto::INT64, dim_);
@@ -92,7 +99,12 @@ void ReduceMeanMapper::Opset11() {
       GetAttr(axis_name_, &dim_);
     }
   } else {
-    GetAttr("axis", &dim_);
+    if (HasAttr(axis_name_)) {
+      GetAttr(axis_name_, &dim_);
+    } else {
+      Assert(TryGetInputValue(axis_name_, &dim_),
+        "Can not get input 'axis(dim)' value.");
+    }
     if (dim_.size() == 0) {
       reduce_all_ = true;
     } else {