未验证 提交 b60ab6b6 编写于 作者: C chentianyu03 提交者: GitHub

change unique op VisitaDataType from small to tiny (#32872)

上级 eeca9639
...@@ -405,13 +405,13 @@ class UniqueKernel : public framework::OpKernel<T> { ...@@ -405,13 +405,13 @@ class UniqueKernel : public framework::OpKernel<T> {
bool return_counts = context.Attr<bool>("return_counts"); bool return_counts = context.Attr<bool>("return_counts");
if (axis_vec.empty()) { if (axis_vec.empty()) {
framework::VisitDataTypeSmall( framework::VisitDataTypeTiny(
data_type, data_type,
UniqueFlattendTensorFunctor<DeviceContext, T>( UniqueFlattendTensorFunctor<DeviceContext, T>(
context, *x, out, return_index, return_inverse, return_counts)); context, *x, out, return_index, return_inverse, return_counts));
} else { } else {
int axis = axis_vec[0]; int axis = axis_vec[0];
framework::VisitDataTypeSmall( framework::VisitDataTypeTiny(
data_type, UniqueDimFunctor<DeviceContext, T>( data_type, UniqueDimFunctor<DeviceContext, T>(
context, *x, out, axis, return_index, return_inverse, context, *x, out, axis, return_index, return_inverse,
return_counts)); return_counts));
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册