[MLU] add int64 support for mlu one_hot_v2 (#46313)

9cc3b28d · Chenxiao Niu · GitHub · 35902c67 · 9cc3b28d
隐藏空白更改
内联并排

Showing with 3 addition and 1 deletion

paddle/fluid/operators/one_hot_v2_op_mlu.cc paddle/fluid/operators/one_hot_v2_op_mlu.cc +3 -1

未找到文件。
--- a/paddle/fluid/operators/one_hot_v2_op_mlu.cc
+++ b/paddle/fluid/operators/one_hot_v2_op_mlu.cc
@@ -97,4 +97,6 @@ class OneHotV2MLUKernel : public framework::OpKernel<T> {
 namespace ops = paddle::operators;
 namespace plat = paddle::platform;

-REGISTER_OP_MLU_KERNEL(one_hot_v2, ops::OneHotV2MLUKernel<int32_t>);
+REGISTER_OP_MLU_KERNEL(one_hot_v2,
+                       ops::OneHotV2MLUKernel<int32_t>,
+                       ops::OneHotV2MLUKernel<int64_t>);