fix_mode_and_bincount_api

PaddlePaddle · Mar 25, 2024 · 2564443 · 2564443
1 parent 4ef550f
commit 2564443
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -279,11 +279,14 @@ void BincountInferMeta(const MetaTensor& x,
   }
   out->set_dims(common::make_ddim({-1}));
   if (weights) {
-    out->set_dtype(weights.dtype());
+    if (weights.dtype() == DataType::FLOAT32) {
+      out->set_dtype(DataType::FLOAT32);
+    } else {
+      out->set_dtype(DataType::FLOAT64);
+    }
   } else {
-    out->set_dtype(x.dtype());
+    out->set_dtype(DataType::INT64);
   }
-
   out->share_lod(x);
 }
 

diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -2537,7 +2537,7 @@ void ModeInferMeta(const MetaTensor& x,
 
   indices->set_dims(dims);
   indices->share_lod(x);
-  indices->set_dtype(x.dtype());
+  indices->set_dtype(DataType::INT64);
 }
 
 void MultinomialInferMeta(const MetaTensor& x,