未验证 提交 4a77cf53 编写于 作者: H houj04 提交者: GitHub

[XPU] add registration of SplitWithNumKernel with int64. (#54478)

上级 24efc13b
...@@ -726,7 +726,8 @@ XPUOpMap& get_kl2_ops() { ...@@ -726,7 +726,8 @@ XPUOpMap& get_kl2_ops() {
{"split_with_num", {"split_with_num",
XPUKernelSet({phi::DataType::FLOAT32, XPUKernelSet({phi::DataType::FLOAT32,
phi::DataType::FLOAT16, phi::DataType::FLOAT16,
phi::DataType::INT32})}, phi::DataType::INT32,
phi::DataType::INT64})},
{"sqrt", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})}, {"sqrt", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
{"sqrt_grad", XPUKernelSet({phi::DataType::FLOAT32})}, {"sqrt_grad", XPUKernelSet({phi::DataType::FLOAT32})},
{"square_grad", {"square_grad",
......
...@@ -79,5 +79,6 @@ PD_REGISTER_KERNEL(split_with_num, ...@@ -79,5 +79,6 @@ PD_REGISTER_KERNEL(split_with_num,
ALL_LAYOUT, ALL_LAYOUT,
phi::SplitWithNumKernel, phi::SplitWithNumKernel,
float, float,
int64_t,
int, int,
phi::dtype::float16) {} phi::dtype::float16) {}
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册