未验证 提交 fc3e51e9 编写于 作者: P PPPPzhang 提交者: GitHub

add output defs for send_u_recv (#51501)

上级 eb0e80a8
...@@ -91,7 +91,6 @@ static std::set<std::string> OpsNeedSetOutputDtypeWhenRegisterPhiKernel = { ...@@ -91,7 +91,6 @@ static std::set<std::string> OpsNeedSetOutputDtypeWhenRegisterPhiKernel = {
"rnn", "rnn",
"search_sort", "search_sort",
"select", "select",
"send_recv",
"send_ue_recv", "send_ue_recv",
"sync_batch_norm_grad", "sync_batch_norm_grad",
"unique", "unique",
......
...@@ -184,4 +184,6 @@ PD_REGISTER_KERNEL(send_u_recv, ...@@ -184,4 +184,6 @@ PD_REGISTER_KERNEL(send_u_recv,
float, float,
double, double,
int, int,
int64_t) {} int64_t) {
kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
}
...@@ -196,4 +196,6 @@ PD_REGISTER_KERNEL(send_u_recv, ...@@ -196,4 +196,6 @@ PD_REGISTER_KERNEL(send_u_recv,
double, double,
int, int,
int64_t, int64_t,
phi::dtype::float16) {} phi::dtype::float16) {
kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
}
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册