[Phi] set where_index out data type

PaddlePaddle · Mar 10, 2022 · 6fd2b1d · 6fd2b1d · paddle-bot-old · Mar 10, 2022
1 parent 83e2e1b
commit 6fd2b1d
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -1268,6 +1268,7 @@ void WhereIndexInferMeta(const MetaTensor& condition, MetaTensor* out) {
       phi::errors::InvalidArgument(
           "Input(Condition) should have number of dimension at least 1"));
   out->set_dims(phi::make_ddim({-1, rank}));
+  out->set_dtype(DataType::INT64);
 }
 
 }  // namespace phi